Fix select_periods dropping metadata, recording, times and empty segments

src/spikeinterface/core/basesorting.py

@@ -891,7 +891,8 @@ def select_periods(self, periods):
         -------
         BaseSorting
             A new sorting object with only samples between start_sample_index and end_sample_index
-            for the given segment_index.
+            for the given segment_index. It has the same segments, properties, annotations, registered
+            recording and times as this sorting.
         """
         from spikeinterface.core.sorting_tools import select_sorting_periods
 

src/spikeinterface/core/numpyextractors.py

@@ -268,9 +268,12 @@ class NumpySorting(BaseSorting):
         The sampling frequency in Hz
     channel_ids : list
         A list of unit_ids.
+    num_segments : int | None, default: None
+        The number of segments. If None, it is inferred from the last spike, which misses
+        trailing segments without spikes.
     """
 
-    def __init__(self, spikes, sampling_frequency, unit_ids):
+    def __init__(self, spikes, sampling_frequency, unit_ids, num_segments=None):
         """ """
         BaseSorting.__init__(self, sampling_frequency, unit_ids)
 
@@ -283,13 +286,20 @@ def __init__(self, spikes, sampling_frequency, unit_ids):
         else:
             nseg = spikes[-1]["segment_index"] + 1
 
+        if num_segments is not None:
+            if spikes.size > 0 and num_segments < nseg:
+                raise ValueError(f"num_segments is {num_segments} but the spikes span {nseg} segments")
+            nseg = num_segments
+
         for segment_index in range(nseg):
             self.add_sorting_segment(SpikeVectorSortingSegment(spikes, segment_index, unit_ids))
 
         # important trick : the cache is already spikes vector
         self._cached_spike_vector = spikes
 
-        self._kwargs = dict(spikes=spikes, sampling_frequency=sampling_frequency, unit_ids=unit_ids)
+        self._kwargs = dict(
+            spikes=spikes, sampling_frequency=sampling_frequency, unit_ids=unit_ids, num_segments=num_segments
+        )
 
     @staticmethod
     def from_sorting(source_sorting: BaseSorting, with_metadata=False, copy_spike_vector=False) -> "NumpySorting":
@@ -300,7 +310,12 @@ def from_sorting(source_sorting: BaseSorting, with_metadata=False, copy_spike_ve
         spike_vector = source_sorting.to_spike_vector()
         if copy_spike_vector:
             spike_vector = spike_vector.copy()
-        sorting = NumpySorting(spike_vector, source_sorting.get_sampling_frequency(), source_sorting.unit_ids.copy())
+        sorting = NumpySorting(
+            spike_vector,
+            source_sorting.get_sampling_frequency(),
+            source_sorting.unit_ids.copy(),
+            num_segments=source_sorting.get_num_segments(),
+        )
         if source_sorting.has_recording():
             sorting._recording = source_sorting._recording
         for source_segment, target_segment in zip(source_sorting._sorting_segments, sorting._sorting_segments):

src/spikeinterface/core/sorting_tools.py

@@ -541,7 +541,8 @@ def select_sorting_periods(sorting: BaseSorting, periods) -> BaseSorting:
     -------
     BaseSorting
         A new sorting object with only samples between start_sample_index and end_sample_index
-        for the given segment_index.
+        for the given segment_index. It has the same segments, properties, annotations, registered
+        recording and times as the input sorting.
     """
     from spikeinterface.core.numpyextractors import NumpySorting
 
@@ -555,11 +556,21 @@ def select_sorting_periods(sorting: BaseSorting, periods) -> BaseSorting:
         sliced_spike_vector = spike_vector[keep_mask]
 
         # important: we keep the original unit ids so the unit_index field in spike vector is still valid
-        sorting = NumpySorting(
-            sliced_spike_vector, sampling_frequency=sorting.sampling_frequency, unit_ids=sorting.unit_ids
+        # we also keep the number of segments, because the last ones can be left without spikes
+        new_sorting = NumpySorting(
+            sliced_spike_vector,
+            sampling_frequency=sorting.sampling_frequency,
+            unit_ids=sorting.unit_ids,
+            num_segments=sorting.get_num_segments(),
         )
-        sorting.copy_metadata(sorting)
-        return sorting
+        sorting.copy_metadata(new_sorting)
+        # not register_recording(): it resets the start times, which would discard a shift_times() of the sorting
+        if sorting.has_recording():
+            new_sorting._recording = sorting._recording
+        for segment, new_segment in zip(sorting.segments, new_sorting.segments):
+            new_segment._t_start = segment._t_start
+            new_segment._native_t_start = segment._native_t_start
+        return new_sorting
     else:
         return sorting
 

src/spikeinterface/core/tests/test_basesorting.py

@@ -17,6 +17,7 @@
     NpzFolderSorting,
     NumpyFolderSorting,
     generate_ground_truth_recording,
+    generate_recording,
     generate_sorting,
     create_sorting_npz,
     generate_sorting,
@@ -439,6 +440,96 @@ def test_select_periods():
     np.testing.assert_array_equal(sliced_sorting.to_spike_vector(), sliced_sorting_array.to_spike_vector())
 
 
+def _make_sorting_for_select_periods(register_recording=True):
+    """Two-segment sorting, with a property, an annotation and a recording that does not start at t = 0."""
+    sampling_frequency = 1000.0
+    sorting = NumpySorting.from_unit_dict(
+        [
+            {"a": np.array([10, 2000, 5000]), "b": np.array([30, 3000])},
+            {"a": np.array([100, 7000]), "b": np.array([40])},
+        ],
+        sampling_frequency=sampling_frequency,
+    )
+    sorting.set_property("quality", np.array(["good", "mua"]))
+    sorting.annotate(probe_name="NP2.0")
+    if register_recording:
+        recording = generate_recording(
+            num_channels=2, sampling_frequency=sampling_frequency, durations=[10.0, 10.0], seed=0
+        )
+        recording.shift_times(500.0, segment_index=0)
+        recording.shift_times(900.0, segment_index=1)
+        sorting.register_recording(recording)
+    return sorting
+
+
+def _spike_times_of_selection(sorting, segment_index):
+    """Spike times of unit "a" after selection, which keeps samples [10, 2000] and [100, 7000] of segments 0 and 1."""
+    periods = np.array([(0, 0, 4000, 0), (1, 0, 8000, 0)], dtype=unit_period_dtype)
+    selected = sorting.select_periods(periods)
+    return selected.get_unit_spike_train("a", segment_index=segment_index, return_times=True)
+
+
+def test_select_periods_keeps_metadata():
+    sorting = _make_sorting_for_select_periods()
+    periods = np.array([(0, 0, 4000, 0), (0, 0, 4000, 1), (1, 0, 8000, 0)], dtype=unit_period_dtype)
+
+    selected = sorting.select_periods(periods)
+
+    assert selected is not sorting
+    np.testing.assert_array_equal(selected.unit_ids, sorting.unit_ids)
+    np.testing.assert_array_equal(selected.get_property("quality"), ["good", "mua"])
+    assert selected.get_annotation("probe_name") == "NP2.0"
+    assert selected.has_recording()
+    assert selected._recording is sorting._recording
+
+
+def test_select_periods_keeps_times():
+    # start times of the recording
+    sorting = _make_sorting_for_select_periods()
+    np.testing.assert_allclose(_spike_times_of_selection(sorting, 0), [500.01, 502.0])
+    np.testing.assert_allclose(_spike_times_of_selection(sorting, 1), [900.1, 907.0])
+
+    # shift of the sorting, on top of the times of the recording
+    sorting.shift_times(2.0, segment_index=1)
+    np.testing.assert_allclose(_spike_times_of_selection(sorting, 0), [500.01, 502.0])
+    np.testing.assert_allclose(_spike_times_of_selection(sorting, 1), [902.1, 909.0])
+
+    # time vector of the recording: irregular, so that a start time cannot stand in for it
+    sorting = _make_sorting_for_select_periods()
+    times = 100.0 + (np.arange(10_000) / 1000.0) ** 2
+    recording = sorting._recording
+    recording.set_times(times, segment_index=0)
+    sorting.register_recording(recording)
+    np.testing.assert_allclose(_spike_times_of_selection(sorting, 0), times[[10, 2000]])
+
+    # shift of the sorting, without any recording
+    sorting = _make_sorting_for_select_periods(register_recording=False)
+    sorting.shift_times(7.0, segment_index=0)
+    np.testing.assert_allclose(_spike_times_of_selection(sorting, 0), [7.01, 9.0])
+    np.testing.assert_allclose(_spike_times_of_selection(sorting, 1), [0.1, 7.0])
+
+
+@pytest.mark.parametrize(
+    "periods, expected_num_spikes",
+    [
+        pytest.param([(0, 0, 4000, 0), (0, 0, 4000, 1)], 4, id="no_period_in_last_segment"),
+        pytest.param([(0, 0, 1, 0)], 0, id="no_spike_in_periods"),
+        pytest.param([], 0, id="no_period"),
+    ],
+)
+def test_select_periods_keeps_segments_without_spikes(periods, expected_num_spikes):
+    sorting = _make_sorting_for_select_periods()
+
+    selected = sorting.select_periods(np.array(periods, dtype=unit_period_dtype))
+
+    assert selected.to_spike_vector().size == expected_num_spikes
+    assert selected.get_num_segments() == sorting.get_num_segments()
+    assert selected.has_recording()
+    assert selected.get_start_time(segment_index=1) == 900.0
+    assert selected.get_unit_spike_train("a", segment_index=1).size == 0
+    assert selected.get_unit_spike_train("a", segment_index=1, return_times=True).size == 0
+
+
 if __name__ == "__main__":
     import tempfile
     from pathlib import Path
@@ -450,3 +541,5 @@ def test_select_periods():
     test_npy_sorting()
     test_empty_sorting()
     test_select_periods()
+    test_select_periods_keeps_metadata()
+    test_select_periods_keeps_times()

src/spikeinterface/core/tests/test_numpy_extractors.py

@@ -120,6 +120,31 @@ def test_NumpySorting(setup_NumpyRecording):
     assert np.shares_memory(sorting2._cached_spike_vector, sorting._cached_spike_vector)
 
 
+def test_NumpySorting_num_segments():
+    sampling_frequency = 30000
+    spikes = np.zeros(2, dtype=minimum_spike_dtype)
+    spikes["sample_index"] = [10, 20]
+
+    # by default, the number of segments is inferred from the last spike
+    assert NumpySorting(spikes, sampling_frequency, ["a"]).get_num_segments() == 1
+
+    # trailing segments without spikes, which must also survive a round trip through the kwargs
+    sorting = NumpySorting(spikes, sampling_frequency, ["a"], num_segments=3)
+    assert sorting.get_num_segments() == 3
+    assert sorting.get_unit_spike_train("a", segment_index=2).size == 0
+    assert load(sorting.to_dict()).get_num_segments() == 3
+    assert NumpySorting.from_sorting(sorting).get_num_segments() == 3
+
+    # no spikes at all
+    no_spikes = np.zeros(0, dtype=minimum_spike_dtype)
+    assert NumpySorting(no_spikes, sampling_frequency, ["a"], num_segments=2).get_num_segments() == 2
+
+    # fewer segments than the spikes span
+    spikes["segment_index"] = [0, 1]
+    with pytest.raises(ValueError, match="num_segments"):
+        NumpySorting(spikes, sampling_frequency, ["a"], num_segments=1)
+
+
 def test_SharedMemorySorting():
     sampling_frequency = 30000
     unit_ids = ["a", "b", "c"]