Use sample metadata 'start_sample' as sample_index_from_session_start

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'
 ]

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:

tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T19-00-00.json

@@ -0,0 +1,3 @@
+{
+    "start_sample": 0
+}

tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T20-00-00.json

@@ -0,0 +1,3 @@
+{
+    "start_sample": 300
+}

tests/resources/chronic-test-data/OnixEphys/OnixEphys_SampleMetadata_2025-05-13T21-00-00.json

@@ -0,0 +1,3 @@
+{
+    "start_sample": 600
+}

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 = (