diff --git a/pyproject.toml b/pyproject.toml index 16a0556..9a7dfca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ dependencies = [ 'aind-data-transformation>=0.0.18', 'spikeinterface[full]>=0.104.0', 'probeinterface>=0.3.2', - 'wavpack-numcodecs>=0.2.2; python_version<"3.13"', + 'wavpack-numcodecs==0.2.2; python_version<"3.13"', 'wavpack-numcodecs>=0.2.3; python_version>="3.13"', 's3fs' ] diff --git a/src/aind_ephys_transformation/ephys_job.py b/src/aind_ephys_transformation/ephys_job.py index ff1b599..58c123d 100644 --- a/src/aind_ephys_transformation/ephys_job.py +++ b/src/aind_ephys_transformation/ephys_job.py @@ -82,6 +82,14 @@ class EphysJobSettings(BasicJobSettings): ), title="Chunks to Compress", ) + chronic_use_sample_metadata: bool = Field( + default=True, + description=( + "If True, it uses the sample metadata to determine the start of " + "sample for each chunk." + ), + title="Chronic Use Sample Metadata Flag", + ) chronic_start_flag: bool = Field( default=False, description=( @@ -288,32 +296,72 @@ def _get_read_blocks(self) -> Iterator[dict]: # noqa: C901 adc_depth = binary_info.pop("adc_depth") recording_list = [] - if self.job_settings.chronic_start_flag: - sample_index_from_session_start = 0 - else: - # If not chronic start flag, we need to parse all previous - # clock bin files to get the cumulative start frame - first_chunk_to_compress = ( - self.job_settings.chronic_chunks_to_compress[0] - ) - first_date_to_compress = datetime.strptime( - first_chunk_to_compress, "%Y-%m-%dT%H-%M-%S" - ) - all_previous_clock_files = [ - p - for p in dataset_folder.glob("**/OnixEphys_Clock_*") - if extract_datetime(p) < first_date_to_compress - ] - sorted_clock_files = sorted( - all_previous_clock_files, - key=lambda x: extract_datetime(x), - ) - sample_index_from_session_start = 0 - for clock_file in sorted_clock_files: - clock_data = np.memmap( - filename=clock_file, dtype="uint64", mode="r" + + # Look for sample metadata files to determine the sample index + # from session start. If not available or invalid, fall back to + # parsing clock files to get cumulative start frame. + are_sample_metadata_files_valid = ( + self._are_sample_metadata_files_valid(onix_folder) + ) + + sample_index_from_session_start = None + if are_sample_metadata_files_valid and ( + self.job_settings.chronic_use_sample_metadata + ): + start_samples = [] + for amplifier_dataset in amplifier_datasets_to_compress: + p = amplifier_dataset + sample_metadata_file = Path(str(p).replace( + "AmplifierData", + "SampleMetadata" + ).replace(".bin", ".json")) + if sample_metadata_file.exists(): + with open(sample_metadata_file) as f: + sample_metadata = json.load(f) + start_sample = sample_metadata.get("start_sample") + # The sample_start should be a non-negative integer. + # This check is in place to spot overflow errors in + # existing datasets. + if start_sample is not None and start_sample >= 0: + start_samples.append(start_sample) + if len(start_samples) == len(amplifier_datasets_to_compress): + # If we have valid start_sample for all datasets, we can + # use it + sample_index_from_session_start = start_samples[0] + logging.info( + "Using start_sample from SampleMetadata files to " + "determine sample index from session start." + ) + + # Fallback to parsing clock files if sample metadata files are + # not valid or some are missing + if sample_index_from_session_start is None: + if self.job_settings.chronic_start_flag: + sample_index_from_session_start = 0 + else: + # If not chronic start flag, we need to parse all previous + # clock bin files to get the cumulative start frame + first_chunk_to_compress = ( + self.job_settings.chronic_chunks_to_compress[0] + ) + first_date_to_compress = datetime.strptime( + first_chunk_to_compress, "%Y-%m-%dT%H-%M-%S" + ) + all_previous_clock_files = [ + p + for p in dataset_folder.glob("**/OnixEphys_Clock_*") + if extract_datetime(p) < first_date_to_compress + ] + sorted_clock_files = sorted( + all_previous_clock_files, + key=lambda x: extract_datetime(x), ) - sample_index_from_session_start += len(clock_data) + sample_index_from_session_start = 0 + for clock_file in sorted_clock_files: + clock_data = np.memmap( + filename=clock_file, dtype="uint64", mode="r" + ) + sample_index_from_session_start += len(clock_data) logging.info( f"Sample index from session start: " f"{sample_index_from_session_start}" @@ -477,6 +525,53 @@ def _get_streams_to_clip(self) -> Iterator[dict]: "n_chan": n_chan, } + def _are_sample_metadata_files_valid(self, onix_folder: Path) -> bool: + """ + Check if sample metadata files are present in the ONIX folder and are + valid (all greater than 0, no overlaps). + Returns an empty list if no valid sample metadata files are found. + + Parameters + ---------- + onix_folder : Path + Path to the ONIX folder. + + Returns + ------- + List[Path] + List of paths to the sample metadata files. + """ + sample_metadata_files = [ + p for p in onix_folder.iterdir() if "SampleMetadata" in p.name + and p.suffix == ".json" + ] + # Sort by date + sample_metadata_files = sorted( + sample_metadata_files, + key=lambda x: extract_datetime(x), + ) + previous_end_sample = -1 + for sample_metadata_file in sample_metadata_files: + with open(sample_metadata_file) as f: + sample_metadata = json.load(f) + start_sample = sample_metadata.get("start_sample") + if start_sample is None or start_sample < 0: # pragma: no cover + logging.warning( + f"Invalid start_sample in {sample_metadata_file}. " + "Expected a non-negative integer. " + "This file will be ignored." + ) + return False + if start_sample <= previous_end_sample: # pragma: no cover + logging.warning( + f"Overlapping start_sample in {sample_metadata_file}. " + "This file will be ignored." + ) + return False + previous_end_sample = start_sample + + return True + def _sync_chronic_timestamps( self, clock_data: np.ndarray, harp_df: pd.DataFrame, fs: float ) -> np.ndarray: diff --git a/tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T19-00-00.json b/tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T19-00-00.json new file mode 100644 index 0000000..46d453a --- /dev/null +++ b/tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T19-00-00.json @@ -0,0 +1,3 @@ +{ + "start_sample": 0 +} \ No newline at end of file diff --git a/tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T20-00-00.json b/tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T20-00-00.json new file mode 100644 index 0000000..9bee514 --- /dev/null +++ b/tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T20-00-00.json @@ -0,0 +1,3 @@ +{ + "start_sample": 300 +} \ No newline at end of file diff --git a/tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T21-00-00.json b/tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T21-00-00.json new file mode 100644 index 0000000..328c68e --- /dev/null +++ b/tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T21-00-00.json @@ -0,0 +1,3 @@ +{ + "start_sample": 600 +} \ No newline at end of file diff --git a/tests/test_ephys_job.py b/tests/test_ephys_job.py index 4549da5..3a21421 100644 --- a/tests/test_ephys_job.py +++ b/tests/test_ephys_job.py @@ -1348,6 +1348,7 @@ def setUpClass(cls): output_directory=Path("output_dir_chronic_append"), compress_job_save_kwargs={"n_jobs": 1}, chronic_chunks_to_compress=["2025-05-13T19-00-00"], + chronic_use_sample_metadata=False, reader_name="chronic", chronic_start_flag=True, ) @@ -1364,6 +1365,7 @@ def setUpClass(cls): "2025-05-13T20-00-00", "2025-05-13T21-00-00", ], + chronic_use_sample_metadata=False, reader_name="chronic", ) cls.chronic_job_settings_append2 = chronic_job_settings_append2 @@ -1422,6 +1424,21 @@ def setUpClass(cls): job_settings=chronic_job_settings_multi_match ) + chronic_job_settings_sampledata = EphysJobSettings( + input_source=CHRONIC_DATA_DIR, + output_directory=Path("output_dir_chronic"), + compress_job_save_kwargs={"n_jobs": 1}, + reader_name="chronic", + chronic_start_flag=False, + chronic_use_sample_metadata=False, + ) + cls.chronic_job_settings_sampledata = ( + chronic_job_settings_sampledata + ) + cls.chronic_job_settings_sampledata = EphysCompressionJob( + job_settings=chronic_job_settings_sampledata + ) + @classmethod def tearDownClass(cls): """Remove output directories created during tests""" @@ -1433,6 +1450,7 @@ def tearDownClass(cls): cls.chronic_job_filter, cls.chronic_job_no_match, cls.chronic_job_multi_match, + cls.chronic_job_settings_sampledata, ]: output_dir = job.job_settings.output_directory if Path(output_dir).exists(): @@ -1495,6 +1513,47 @@ def test_get_read_blocks_filter(self): read_blocks_repr_str = set([json.dumps(o) for o in read_blocks_repr]) self.assertEqual(expected_scaled_read_blocks_str, read_blocks_repr_str) + def test_read_blocks_sampledata(self): + """Tests _get_read_blocks method when there is sample data in the + metadata""" + chunks = [ + "2025-05-13T19-00-00", + "2025-05-13T20-00-00", + "2025-05-13T21-00-00" + ] + for i, chunk in enumerate(chunks): + self.chronic_job_settings_sampledata.job_settings.\ + chronic_chunks_to_compress = [chunk] + read_blocks = ( + self.chronic_job_settings_sampledata._get_read_blocks() + ) + # In this case, start samples should be 100 samples apart + # since the ephys and clock data have 100 samples per chunk + for read_block in read_blocks: + recording = read_block["recording"] + start_sample = recording.get_annotation( + "sample_index_from_session_start" + ) + self.assertEqual(start_sample, i * 100) + + # Test with sample metadata + self.chronic_job_settings_sampledata.job_settings.\ + chronic_use_sample_metadata = True + for i, chunk in enumerate(chunks): + self.chronic_job_settings_sampledata.job_settings.\ + chronic_chunks_to_compress = [chunk] + read_blocks = ( + self.chronic_job_settings_sampledata._get_read_blocks() + ) + # SampleMetadata has jumps of 300 samples between chunks, + # so start samples should be 300 samples apart + for read_block in read_blocks: + recording = read_block["recording"] + start_sample = recording.get_annotation( + "sample_index_from_session_start" + ) + self.assertEqual(start_sample, i * 300) + def test_read_blocks_no_match(self): """Tests _get_read_blocks method with no matching chunks""" with self.assertRaises(ValueError): @@ -1797,7 +1856,7 @@ def test_s3_location( mock_copy_file_to_s3.assert_has_calls( expected_copy_calls, any_order=True ) - self.assertEqual(mock_copy_file_to_s3.call_count, 8) + self.assertEqual(mock_copy_file_to_s3.call_count, 11) # Assert call to write_or_append_recording_to_zarr expected_zarr_s3_path = (