Merge pull request #3677 from chrishalcrow/add-sorting-time-slice

Add `time_slice` method to `BaseSorting`

Author: alejoe91

src/spikeinterface/core/baserecording.py

@@ -772,7 +772,7 @@ def frame_slice(self, start_frame: int | None, end_frame: int | None) -> BaseRec
 
     def time_slice(self, start_time: float | None, end_time: float) -> BaseRecording:
         """
-        Returns a new recording with sliced time. Note that this operation is not in place.
+        Returns a new recording object, restricted to the time interval [start_time, end_time].
 
         Parameters
         ----------

src/spikeinterface/core/basesorting.py

@@ -179,7 +179,9 @@ def get_unit_spike_train(
             return spike_frames
 
     def register_recording(self, recording, check_spike_frames=True):
-        """Register a recording to the sorting.
+        """
+        Register a recording to the sorting. If the sorting and recording both contain
+        time information, the recording’s time information will be used.
 
         Parameters
         ----------
@@ -472,6 +474,43 @@ def frame_slice(self, start_frame, end_frame, check_spike_frames=True):
         )
         return sub_sorting
 
+    def time_slice(self, start_time: float | None, end_time: float | None) -> BaseSorting:
+        """
+        Returns a new sorting object, restricted to the time interval [start_time, end_time].
+
+        Parameters
+        ----------
+        start_time : float | None, default: None
+            The start time in seconds. If not provided it is set to 0.
+        end_time : float | None, default: None
+            The end time in seconds. If not provided it is set to the total duration.
+
+        Returns
+        -------
+        BaseSorting
+            A new sorting object with only samples between start_time and end_time
+        """
+
+        assert self.get_num_segments() == 1, "Time slicing is only supported for single segment sortings."
+
+        start_frame = self.time_to_sample_index(start_time, segment_index=0) if start_time else None
+        end_frame = self.time_to_sample_index(end_time, segment_index=0) if end_time else None
+
+        return self.frame_slice(start_frame=start_frame, end_frame=end_frame)
+
+    def time_to_sample_index(self, time, segment_index=0):
+        """
+        Transform time in seconds into sample index
+        """
+        if self.has_recording():
+            sample_index = self._recording.time_to_sample_index(time, segment_index=segment_index)
+        else:
+            segment = self._sorting_segments[segment_index]
+            t_start = segment._t_start if segment._t_start is not None else 0
+            sample_index = round((time - t_start) * self.get_sampling_frequency())
+
+        return sample_index
+
     def get_all_spike_trains(self, outputs="unit_id"):
         """
         Return all spike trains concatenated.

src/spikeinterface/core/tests/test_basesorting.py

@@ -25,6 +25,8 @@
 from spikeinterface.core.testing import check_sorted_arrays_equal, check_sortings_equal
 from spikeinterface.core.generate import generate_sorting
 
+from spikeinterface.core import generate_recording, generate_ground_truth_recording
+
 
 def test_BaseSorting(create_cache_folder):
     cache_folder = create_cache_folder
@@ -206,6 +208,32 @@ def test_empty_sorting():
     assert spikes.shape == (0,)
 
 
+def test_time_slice():
+
+    sampling_frequency = 10_000.0
+
+    # no recording attached to sorting
+    sorting = generate_sorting(durations=[1.0], sampling_frequency=sampling_frequency)
+
+    sliced_sorting_times = sorting.time_slice(start_time=0.1, end_time=0.8)
+    sliced_sorting_frames = sorting.frame_slice(start_frame=1000, end_frame=8000)
+
+    assert np.allclose(
+        sliced_sorting_times.to_spike_vector()["sample_index"], sliced_sorting_frames.to_spike_vector()["sample_index"]
+    )
+
+    # with recording
+    recording, sorting = generate_ground_truth_recording(durations=[1.0], sampling_frequency=sampling_frequency)
+    sorting.register_recording(recording)
+
+    sliced_sorting_times = sorting.time_slice(start_time=0.1, end_time=0.8)
+    sliced_sorting_frames = sorting.frame_slice(start_frame=1000, end_frame=8000)
+
+    assert np.allclose(
+        sliced_sorting_times.to_spike_vector()["sample_index"], sliced_sorting_frames.to_spike_vector()["sample_index"]
+    )
+
+
 if __name__ == "__main__":
     test_BaseSorting()
     test_npy_sorting()