Repository: SpikeInterface/spikeinterface
Pull request #1425 (2023-03-16 15:47:35)

Change default value for end_frame in FrameSliceSorting
Hi
In FrameSliceSorting:
If None (default), the end_frame is now either:
    - The total number of samples, if a recording is assigned to the sorting
    - The maximum spike frame + 1, otherwise

This ensures that `Sorting.frame_slice(start_frame=None, end_frame=None)` not only 
returns all the spikes (as previously), but also that the sortings' registered recording remains complete.

Note that the following configs are now legal (and return empty sortings)
- `start_frame < end_frame <= min_spike_time`
- `max_spike_time < start_frame < end_frame`


Patch:
diff --git a/spikeinterface/core/frameslicerecording.py b/spikeinterface/core/frameslicerecording.py
index 14e23ee1b..8cc442029 100644
--- a/spikeinterface/core/frameslicerecording.py
+++ b/spikeinterface/core/frameslicerecording.py
@@ -10,6 +10,18 @@ class FrameSliceRecording(BaseRecording):
 
     Do not use this class directly but use `recording.frame_slice(...)`
 
+    Parameters
+    ----------
+    parent_recording: BaseRecording
+    start_frame: None or int
+        Earliest included frame in the parent recording.
+        Times are re-referenced to start_frame in the
+        sliced object. Set to 0 by default.
+    end_frame: None or int
+        Latest frame in the parent recording. As for usual
+        python slicing, the end frame is excluded. 
+        Set to the recording's total number of samples by 
+        default
     """
 
     def __init__(self, parent_recording, start_frame=None, end_frame=None):
diff --git a/spikeinterface/core/frameslicesorting.py b/spikeinterface/core/frameslicesorting.py
index 8f0be4320..d5e2b3a28 100644
--- a/spikeinterface/core/frameslicesorting.py
+++ b/spikeinterface/core/frameslicesorting.py
@@ -2,6 +2,7 @@ import numpy as np
 import warnings
 
 from .basesorting import BaseSorting, BaseSortingSegment
+from .waveform_tools import has_exceeding_spikes
 
 
 class FrameSliceSorting(BaseSorting):
@@ -11,6 +12,25 @@ class FrameSliceSorting(BaseSorting):
 
     Do not use this class directly but use `sorting.frame_slice(...)`
 
+    When a recording is registered for the parent sorting,
+    a corresponding sliced recording is registered to the sliced sorting.
+
+    Note that the returned sliced sorting may be empty.
+
+    Parameters
+    ----------
+    parent_sorting: BaseSorting
+    start_frame: None or int
+        Earliest included frame in the parent sorting(/recording).
+        Spike times(/traces) are re-referenced to start_frame in the
+        sliced objects. Set to 0 by default.
+    end_frame: None or int
+        Latest frame in the parent sorting(/recording). As for usual
+        python slicing, the end frame is excluded (such that the max
+        spike frame in the sliced sorting is `end_frame - start_frame - 1`)
+        If None (default), the end_frame is either:
+            - The total number of samples, if a recording is assigned
+            - The maximum spike frame + 1, if no recording is assigned
     """
 
     def __init__(self, parent_sorting, start_frame=None, end_frame=None):
@@ -18,20 +38,41 @@ class FrameSliceSorting(BaseSorting):
 
         assert parent_sorting.get_num_segments() == 1, 'FrameSliceSorting work only with one segment'
 
-        if start_frame is not None or end_frame is None:
-            parent_size = 0
-            for u in parent_sorting.get_unit_ids():
-                parent_size = np.max([parent_size, np.max(parent_sorting.get_unit_spike_train(u))])
 
         if start_frame is None:
             start_frame = 0
-        else:
-            assert 0 <= start_frame < parent_size
+        assert 0 <= start_frame, "Invalid value for start_frame: expected positive integer."
 
-        if end_frame is None:
-            end_frame = parent_size + 1
+        if parent_sorting.has_recording():
+            # Pull df end_frame from recording
+            parent_n_samples = parent_sorting._recording.get_total_samples()
+            if end_frame is None:
+                end_frame = parent_n_samples
+            assert end_frame <= parent_n_samples, (
+                "`end_frame` should be smaller than the sortings total number of samples."
+            )
+            assert start_frame <= parent_n_samples, (
+                "`start_frame` should be smaller than the sortings total number of samples."
+            )
+            if has_exceeding_spikes(parent_sorting._recording, parent_sorting):
+                raise ValueError(
+                    "The sorting object has spikes exceeding the recording duration. You have to remove those spikes "
+                    "with the `spikeinterface.curation.remove_excess_spikes()` function"
+                )
         else:
-            assert end_frame > start_frame, "'start_frame' must be smaller than 'end_frame'!"
+            # Pull df end_frame from spikes
+            if end_frame is None:
+                max_spike_time = 0
+                for u in parent_sorting.get_unit_ids():
+                    max_spike_time = np.max([max_spike_time, np.max(parent_sorting.get_unit_spike_train(u))])
+                end_frame = max_spike_time + 1
+
+        assert start_frame < end_frame, (
+            "`start_frame` should be greater than `end_frame`. "
+            "This may be due to start_frame >= max_spike_time, if the end frame "
+            "was not specified explicitly."
+        )
+            
 
         BaseSorting.__init__(self,
                              sampling_frequency=parent_sorting.get_sampling_frequency(),


Test patch:
diff --git a/spikeinterface/core/tests/test_frameslicesorting.py b/spikeinterface/core/tests/test_frameslicesorting.py
index d64b0a4fe..1e4b46916 100644
--- a/spikeinterface/core/tests/test_frameslicesorting.py
+++ b/spikeinterface/core/tests/test_frameslicesorting.py
@@ -1,32 +1,111 @@
-from spikeinterface.core import generate_sorting
+import warnings
+
+import numpy as np
+from numpy.testing import assert_raises
+
+from spikeinterface.core import NumpyRecording, NumpySorting
 
 
 def test_FrameSliceSorting():
-    fs = 30000
-    duration = 10
-    sort = generate_sorting(num_units=10,  durations=[
-                            duration], sampling_frequency=fs)
 
-    mid_frame = (duration * fs) // 2
+    # Single segment sorting, with and without attached recording
+    # Since the default end_frame can be set either from the last spike
+    # or from the registered recording
+    sf = 10
+    nsamp = 1000
+    max_spike_time = 900
+    min_spike_time = 100
+    unit_0_train = np.arange(min_spike_time + 10, max_spike_time - 10)
+    spike_times = {
+        "0": unit_0_train,
+        "1": np.arange(min_spike_time, max_spike_time),
+    }
+    # Sorting with attached rec
+    sorting = NumpySorting.from_dict( [spike_times], sf)
+    rec = NumpyRecording([np.zeros((nsamp, 5))], sampling_frequency=sf)
+    sorting.register_recording(rec)
+    # Sorting without attached rec
+    sorting_norec = NumpySorting.from_dict( [spike_times], sf)
+    # Sorting with attached rec and exceeding spikes
+    sorting_exceeding = NumpySorting.from_dict( [spike_times], sf)
+    rec_exceeding = NumpyRecording([np.zeros((max_spike_time-1, 5))], sampling_frequency=sf)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore")
+        sorting_exceeding.register_recording(rec_exceeding)
+
+    mid_frame = nsamp // 2
+
     # duration of all slices is mid_frame. Spike trains are re-referenced to the start_time
-    sub_sort = sort.frame_slice(None, None)
-    for u in sort.get_unit_ids():
-        assert len(sort.get_unit_spike_train(u)) == len(
-            sub_sort.get_unit_spike_train(u))
+    # Vary start_frame/end_frame combination
+    start_frame, end_frame = None, None
+    sub_sorting = sorting.frame_slice(start_frame, end_frame)
+    assert np.array_equal(sub_sorting.get_unit_spike_train("0"), unit_0_train)
+    assert sub_sorting._recording.get_total_samples() == nsamp
+    sub_sorting_norec = sorting.frame_slice(start_frame, end_frame)
+    assert np.array_equal(sub_sorting_norec.get_unit_spike_train("0"), unit_0_train)
+
+    start_frame, end_frame = None, mid_frame
+    sub_sorting = sorting.frame_slice(start_frame, end_frame)
+    assert np.array_equal(
+        sub_sorting.get_unit_spike_train("0"), 
+        [t for t in unit_0_train if t < mid_frame]
+    )
+    assert sub_sorting._recording.get_total_samples() == mid_frame
+    sub_sorting_norec = sorting.frame_slice(start_frame, end_frame)
+    assert np.array_equal(
+        sub_sorting_norec.get_unit_spike_train("0"), 
+        sub_sorting.get_unit_spike_train("0")
+    )
+
+    start_frame, end_frame = mid_frame, None
+    sub_sorting = sorting.frame_slice(start_frame, end_frame)
+    assert np.array_equal(
+        sub_sorting.get_unit_spike_train("0"), 
+        [t - mid_frame for t in unit_0_train if t >= mid_frame]
+    )
+    assert sub_sorting._recording.get_total_samples() == nsamp - mid_frame
+    sub_sorting_norec = sorting.frame_slice(start_frame, end_frame)
+    assert np.array_equal(
+        sub_sorting_norec.get_unit_spike_train("0"), 
+        sub_sorting.get_unit_spike_train("0")
+    )
+
+    start_frame, end_frame = mid_frame - 10, mid_frame + 10
+    sub_sorting = sorting.frame_slice(start_frame, end_frame)
+    assert np.array_equal(
+        sub_sorting.get_unit_spike_train("0"), 
+        [t - start_frame for t in unit_0_train if start_frame <= t < end_frame]
+    )
+    assert sub_sorting._recording.get_total_samples() == 20
+    sub_sorting_norec = sorting.frame_slice(start_frame, end_frame)
+    assert np.array_equal(
+        sub_sorting_norec.get_unit_spike_train("0"), 
+        sub_sorting.get_unit_spike_train("0")
+    )
+
+    # Edge cases: start_frame > end_frame
+    assert_raises(Exception, sorting.frame_slice, 100, 90)
 
-    sub_sort = sort.frame_slice(None, mid_frame)
-    for u in sort.get_unit_ids():
-        assert max(sub_sort.get_unit_spike_train(u)) <= mid_frame
+    # Edge case: start_frame > max_spike_time
+    # Fails without rec (since end_frame is last spike)
+    assert_raises(Exception, sorting_norec.frame_slice, max_spike_time + 1, None)
+    # Empty sorting with rec
+    sub_sorting = sorting.frame_slice(max_spike_time + 1, None)
+    assert np.array_equal(
+        sub_sorting.get_unit_spike_train("1"), 
+        []
+    )
 
-    sub_sort = sort.frame_slice(mid_frame, None)
-    for u in sort.get_unit_ids():
-        assert max(sub_sort.get_unit_spike_train(u)) <= mid_frame
+    # Edge case: end_frame <= min_spike_time
+    # Empty sorting 
+    sub_sorting = sorting.frame_slice(None, min_spike_time)
+    assert np.array_equal( sub_sorting.get_unit_spike_train("1"), [])
 
-    sub_sort = sort.frame_slice(
-        mid_frame - mid_frame // 2, mid_frame + mid_frame // 2)
-    for u in sort.get_unit_ids():
-        assert max(sub_sort.get_unit_spike_train(u)) <= mid_frame
+    # Edge case: start_frame = end_frame
+    assert_raises(Exception, sorting.frame_slice, max_spike_time, max_spike_time)
 
+    # Sorting with exceeding spikes
+    assert_raises(Exception, sorting_exceeding.frame_slice, None, None)
 
 if __name__ == '__main__':
     test_FrameSliceSorting()