Repository: SpikeInterface/spikeinterface
Pull request #2353 (2023-12-19 10:05:39)

Cache spike_vector from parent


Patch:
diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py
index ba4f4b485..068165529 100644
--- a/src/spikeinterface/core/basesorting.py
+++ b/src/spikeinterface/core/basesorting.py
@@ -513,6 +513,16 @@ class BaseSorting(BaseExtractor):
                 for unit_id in unit_ids:
                     self.get_unit_spike_train(unit_id, segment_index=segment_index, use_cache=True)
 
+    def _custom_cache_spike_vector(self) -> None:
+        """
+        Function that can be implemented by some children sorting to quickly
+        cache the spike vector without computing it from spike trains
+        (e.g. computing it from a sorting parent).
+        This function should set the `self._cached_spike_vector`, see for
+        instance the `UnitsSelectionSorting` implementation.
+        """
+        pass
+
     def to_spike_vector(
         self, concatenated=True, extremum_channel_inds=None, use_cache=True
     ) -> np.ndarray | list[np.ndarray]:
@@ -548,6 +558,9 @@ class BaseSorting(BaseExtractor):
             spike_dtype = spike_dtype + [("channel_index", "int64")]
             ext_channel_inds = np.array([extremum_channel_inds[unit_id] for unit_id in self.unit_ids])
 
+        if use_cache and self._cached_spike_vector is None:
+            self._custom_cache_spike_vector()
+
         if use_cache and self._cached_spike_vector is not None:
             # the cache already exists
             if extremum_channel_inds is None:
diff --git a/src/spikeinterface/core/unitsselectionsorting.py b/src/spikeinterface/core/unitsselectionsorting.py
index 26f7f1a3b..ed199b06a 100644
--- a/src/spikeinterface/core/unitsselectionsorting.py
+++ b/src/spikeinterface/core/unitsselectionsorting.py
@@ -46,6 +46,25 @@ class UnitsSelectionSorting(BaseSorting):
 
         self._kwargs = dict(parent_sorting=parent_sorting, unit_ids=unit_ids, renamed_unit_ids=renamed_unit_ids)
 
+    def _custom_cache_spike_vector(self) -> None:
+        if self._parent_sorting._cached_spike_vector is None:
+            self._parent_sorting._custom_cache_spike_vector()
+
+            if self._parent_sorting._cached_spike_vector is None:
+                return
+
+        parent_spike_vector = self._parent_sorting._cached_spike_vector
+        parent_unit_indices = self._parent_sorting.ids_to_indices(self._unit_ids)
+        mask = np.isin(parent_spike_vector["unit_index"], parent_unit_indices)
+        spike_vector = np.array(
+            parent_spike_vector[mask]
+        )  # np.array() necessary to fix 'read-only' crash with memmaps.
+        spike_vector["unit_index"] = np.searchsorted(
+            parent_unit_indices, spike_vector["unit_index"]
+        )  # Trick to make sure that the new indices are correct.
+
+        self._cached_spike_vector = spike_vector
+
 
 class UnitsSelectionSortingSegment(BaseSortingSegment):
     def __init__(self, parent_segment, ids_conversion):
diff --git a/src/spikeinterface/curation/remove_excess_spikes.py b/src/spikeinterface/curation/remove_excess_spikes.py
index 399ea6954..76cf2b7f3 100644
--- a/src/spikeinterface/curation/remove_excess_spikes.py
+++ b/src/spikeinterface/curation/remove_excess_spikes.py
@@ -30,10 +30,14 @@ class RemoveExcessSpikesSorting(BaseSorting):
             sorting.get_num_segments() == recording.get_num_segments()
         ), "The sorting and recording objects must have the same number of samples!"
 
+        self._parent_sorting = sorting
+        self._num_samples = np.empty(sorting.get_num_segments(), dtype=np.int64)
         for segment_index in range(sorting.get_num_segments()):
             sorting_segment = sorting._sorting_segments[segment_index]
-            num_samples = recording.get_num_samples(segment_index=segment_index)
-            self.add_sorting_segment(RemoveExcessSpikesSortingSegment(sorting_segment, num_samples))
+            self._num_samples[segment_index] = recording.get_num_samples(segment_index=segment_index)
+            self.add_sorting_segment(
+                RemoveExcessSpikesSortingSegment(sorting_segment, self._num_samples[segment_index])
+            )
 
         sorting.copy_metadata(self, only_main=False)
         if sorting.has_recording():
@@ -41,6 +45,27 @@ class RemoveExcessSpikesSorting(BaseSorting):
 
         self._kwargs = {"sorting": sorting, "recording": recording}
 
+    def _custom_cache_spike_vector(self) -> None:
+        if self._parent_sorting._cached_spike_vector is None:
+            self._parent_sorting._custom_cache_spike_vector()
+
+            if self._parent_sorting._cached_spike_vector is None:
+                return
+
+        parent_spike_vector = self._parent_sorting._cached_spike_vector
+
+        list_spike_vectors = []
+        segments_bounds = np.searchsorted(
+            parent_spike_vector["segment_index"], np.arange(1 + parent_spike_vector["segment_index"][-1])
+        )
+        for segment_index in range(parent_spike_vector["segment_index"][-1]):
+            spike_vector = parent_spike_vector[segments_bounds[segment_index] : segments_bounds[segment_index + 1]]
+            end = np.searchsorted(spike_vector["sample_index"], self._num_samples[segment_index])
+            list_spike_vectors.append(spike_vector[:end])
+
+        spike_vector = np.concatenate(list_spike_vectors)
+        self._cached_spike_vector = spike_vector
+
 
 class RemoveExcessSpikesSortingSegment(BaseSortingSegment):
     def __init__(self, parent_segment: BaseSortingSegment, num_samples: int) -> None:


Test patch:
diff --git a/src/spikeinterface/core/tests/test_unitsselectionsorting.py b/src/spikeinterface/core/tests/test_unitsselectionsorting.py
index 4c7d1746d..4ffc370da 100644
--- a/src/spikeinterface/core/tests/test_unitsselectionsorting.py
+++ b/src/spikeinterface/core/tests/test_unitsselectionsorting.py
@@ -38,5 +38,16 @@ def test_failure_with_non_unique_unit_ids():
         sorting2 = UnitsSelectionSorting(sorting, unit_ids=[0, 2], renamed_unit_ids=["a", "a"])
 
 
+def test_custom_cache_spike_vector():
+    sorting = generate_sorting(num_units=3, durations=[0.100, 0.100], sampling_frequency=30000.0)
+
+    sub_sorting = UnitsSelectionSorting(sorting, unit_ids=[0, 2], renamed_unit_ids=["b", "a"])
+    spike_vector = sub_sorting.to_spike_vector(use_cache=False)
+
+    sorting.to_spike_vector(use_cache=True)
+    sub_sorting._custom_cache_spike_vector()
+    assert np.all(spike_vector == sub_sorting._cached_spike_vector)
+
+
 if __name__ == "__main__":
     test_basic_functions()