Repository: SpikeInterface/spikeinterface
Pull request #3677 (2025-02-07 12:29:07)

Add `time_slice` method to `BaseSorting`
In response to #3675

This PR allows the user to slice a sorting, so that you can do something like

``` python
reduced_sorting = sorting.time_slice(start_time = 60, end_time=120)
reduced_spike_train = reduced_sorting.to_spike_vector()
```

Code and test were essentially copied and pasted from the `BaseRecording` object's `time_slice` method. Then had to make a new `time_to_sample_index` function for the `BaseSorting`, whose logic was taken from `get_unit_spike_trains` here: 
https://github.com/SpikeInterface/spikeinterface/blob/4e34fc747115e91dde06bf2b42ea6bfe4d42616c/src/spikeinterface/core/basesorting.py#L169
This uses the attached recording if there is one, checks if there's a `_t_start`, and sets `_t_start` to zero if there isn't.

Only works for single-segment sorting objects, which matches the capabilities of `time_slice` for the `BaseRecording` object.


Patch:
diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py
index 3e7283090..d1e44f445 100644
--- a/src/spikeinterface/core/baserecording.py
+++ b/src/spikeinterface/core/baserecording.py
@@ -772,7 +772,7 @@ class BaseRecording(BaseRecordingSnippets):
 
     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
         ----------
diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py
index 9a0e242d6..23f1ce715 100644
--- a/src/spikeinterface/core/basesorting.py
+++ b/src/spikeinterface/core/basesorting.py
@@ -179,7 +179,9 @@ class BaseSorting(BaseExtractor):
             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 @@ class BaseSorting(BaseExtractor):
         )
         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.


Test patch:
diff --git a/src/spikeinterface/core/tests/test_basesorting.py b/src/spikeinterface/core/tests/test_basesorting.py
index 557617ae1..c050a0237 100644
--- a/src/spikeinterface/core/tests/test_basesorting.py
+++ b/src/spikeinterface/core/tests/test_basesorting.py
@@ -25,6 +25,8 @@ from spikeinterface.core.base import BaseExtractor
 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()