Rename the spike vector caches to state which order they hold

src/spikeinterface/core/generate.py

@@ -365,18 +365,18 @@ def __init__(
         BaseSorting.__init__(self, sampling_frequency, unit_ids)
 
         self.parent_unit_ids = sorting.unit_ids
-        self._cached_spike_vector = sorting.to_spike_vector().copy()
+        spikes = sorting.to_spike_vector().copy()
         self.refractory_period_ms = refractory_period_ms
 
-        self.added_spikes_from_existing_mask = np.zeros(len(self._cached_spike_vector), dtype=bool)
-        self.added_spikes_from_new_mask = np.zeros(len(self._cached_spike_vector), dtype=bool)
+        self.added_spikes_from_existing_mask = np.zeros(len(spikes), dtype=bool)
+        self.added_spikes_from_new_mask = np.zeros(len(spikes), dtype=bool)
 
         if added_spikes_existing_units is not None and len(added_spikes_existing_units) > 0:
             assert (
                 added_spikes_existing_units.dtype == minimum_spike_dtype
             ), "added_spikes_existing_units should be a spike vector"
             added_unit_indices = np.arange(len(self.parent_unit_ids))
-            self._cached_spike_vector = np.concatenate((self._cached_spike_vector, added_spikes_existing_units))
+            spikes = np.concatenate((spikes, added_spikes_existing_units))
             self.added_spikes_from_existing_mask = np.concatenate(
                 (self.added_spikes_from_existing_mask, np.ones(len(added_spikes_existing_units), dtype=bool))
             )
@@ -388,22 +388,28 @@ def __init__(
             assert (
                 added_spikes_new_units.dtype == minimum_spike_dtype
             ), "added_spikes_new_units should be a spike vector"
-            self._cached_spike_vector = np.concatenate((self._cached_spike_vector, added_spikes_new_units))
+            spikes = np.concatenate((spikes, added_spikes_new_units))
             self.added_spikes_from_existing_mask = np.concatenate(
                 (self.added_spikes_from_existing_mask, np.zeros(len(added_spikes_new_units), dtype=bool))
             )
             self.added_spikes_from_new_mask = np.concatenate(
                 (self.added_spikes_from_new_mask, np.ones(len(added_spikes_new_units), dtype=bool))
             )
 
-        sort_idxs = np.lexsort([self._cached_spike_vector["sample_index"], self._cached_spike_vector["segment_index"]])
-        self._cached_spike_vector = self._cached_spike_vector[sort_idxs]
+        sort_idxs = np.lexsort(
+            [
+                spikes["sample_index"],
+                spikes["segment_index"],
+            ]
+        )
+        spikes = spikes[sort_idxs]
+        self._cached_time_ordered = {"spikes": spikes, "segment_slices": None}
         self.added_spikes_from_existing_mask = self.added_spikes_from_existing_mask[sort_idxs]
         self.added_spikes_from_new_mask = self.added_spikes_from_new_mask[sort_idxs]
 
         # We need to add the sorting segments
         for segment_index in range(sorting.get_num_segments()):
-            segment = SpikeVectorSortingSegment(self._cached_spike_vector, segment_index, unit_ids=self.unit_ids)
+            segment = SpikeVectorSortingSegment(spikes, segment_index, unit_ids=self.unit_ids)
             self.add_sorting_segment(segment)
 
         if self.refractory_period_ms is not None:
@@ -556,20 +562,18 @@ def clean_refractory_period(self):
         ## This function will remove the added spikes that will violate RPV, but does not affect the
         ## spikes in the original sorting. So if some RPV violation are present in this sorting,
         ## they will be left untouched
-        unit_indices = np.unique(self._cached_spike_vector["unit_index"])
+        spikes = self._cached_time_ordered["spikes"]
+        unit_indices = np.unique(spikes["unit_index"])
         rpv = int(self.get_sampling_frequency() * self.refractory_period_ms / 1000)
         to_keep = ~self.added_spikes_from_existing_mask.copy()
         for segment_index in range(self.get_num_segments()):
             for unit_ind in unit_indices:
-                (indices,) = np.nonzero(
-                    (self._cached_spike_vector["unit_index"] == unit_ind)
-                    * (self._cached_spike_vector["segment_index"] == segment_index)
-                )
+                (indices,) = np.nonzero((spikes["unit_index"] == unit_ind) * (spikes["segment_index"] == segment_index))
                 to_keep[indices[1:]] = np.logical_or(
-                    to_keep[indices[1:]], np.diff(self._cached_spike_vector[indices]["sample_index"]) > rpv
+                    to_keep[indices[1:]], np.diff(spikes[indices]["sample_index"]) > rpv
                 )
 
-        self._cached_spike_vector = self._cached_spike_vector[to_keep]
+        self._cached_time_ordered = {"spikes": spikes[to_keep], "segment_slices": None}
         self.added_spikes_from_existing_mask = self.added_spikes_from_existing_mask[to_keep]
         self.added_spikes_from_new_mask = self.added_spikes_from_new_mask[to_keep]
 

src/spikeinterface/core/numpyextractors.py

@@ -287,7 +287,7 @@ def __init__(self, spikes, sampling_frequency, unit_ids):
             self.add_sorting_segment(SpikeVectorSortingSegment(spikes, segment_index, unit_ids))
 
         # important trick : the cache is already spikes vector
-        self._cached_spike_vector = spikes
+        self._cached_time_ordered = {"spikes": spikes, "segment_slices": None}
 
         self._kwargs = dict(spikes=spikes, sampling_frequency=sampling_frequency, unit_ids=unit_ids)
 
@@ -523,7 +523,7 @@ def __init__(self, shm_name, shape, sampling_frequency, unit_ids, dtype=minimum_
             self.add_sorting_segment(SpikeVectorSortingSegment(self.shm_spikes, segment_index, unit_ids))
 
         # important trick : the cache is already spikes vector
-        self._cached_spike_vector = self.shm_spikes
+        self._cached_time_ordered = {"spikes": self.shm_spikes, "segment_slices": None}
 
         # this is very important for the shm.unlink()
         # only the main instance need to call it

src/spikeinterface/core/sortingfolder.py

@@ -53,7 +53,7 @@ def __init__(self, folder_path, mmap_mode: str | None = None):
         for segment_index in range(num_segments):
             self.add_sorting_segment(SpikeVectorSortingSegment(self.spikes, segment_index, unit_ids))
         # important trick : the cache is already spikes vector
-        self._cached_spike_vector = self.spikes
+        self._cached_time_ordered = {"spikes": self.spikes, "segment_slices": None}
 
         load_properties_from_folder(folder_path / "properties", self)
         load_annotations_from_folder(folder_path, self)

src/spikeinterface/core/tests/test_basesorting.py

@@ -122,7 +122,7 @@ def test_BaseSorting(create_cache_folder):
 
     spikes = sorting.to_spike_vector()
     # print(spikes)
-    assert sorting._cached_spike_vector is not None
+    assert sorting._cached_time_ordered is not None
 
     spikes = sorting.to_spike_vector()
     ordered_spikes, order, slices = sorting.to_reordered_spike_vector(
@@ -132,17 +132,6 @@ def test_BaseSorting(create_cache_folder):
     )
     assert np.array_equal(spikes[order], ordered_spikes)
 
-    ordered_spikes, order, slices = sorting.to_reordered_spike_vector(
-        lexsort=(
-            "sample_index",
-            "unit_index",
-            "segment_index",
-        ),
-        return_order=True,
-        return_slices=True,
-    )
-    assert np.array_equal(spikes[order], ordered_spikes)
-
     num_spikes_per_unit = sorting.count_num_spikes_per_unit(outputs="dict")
     num_spikes_per_unit = sorting.count_num_spikes_per_unit(outputs="array")
     total_spikes = sorting.count_total_num_spikes()
@@ -211,11 +200,7 @@ def test_zarr_save_with_sharding(create_cache_folder):
 
 
 @pytest.mark.parametrize("use_numba", [True, False], ids=["numba", "numpy"])
-@pytest.mark.parametrize(
-    "lexsort",
-    [("sample_index", "segment_index", "unit_index"), ("sample_index", "unit_index", "segment_index")],
-)
-def test_to_reordered_spike_vector(lexsort, use_numba, monkeypatch):
+def test_to_reordered_spike_vector(use_numba, monkeypatch):
     """`to_reordered_spike_vector` should group spikes by (unit, segment) without disturbing them."""
     if use_numba and importlib.util.find_spec("numba") is None:
         pytest.skip("numba not installed")
@@ -231,17 +216,12 @@ def test_to_reordered_spike_vector(lexsort, use_numba, monkeypatch):
     )
 
     ordered_spikes, order, slices = sorting.to_reordered_spike_vector(
-        lexsort=lexsort, return_order=True, return_slices=True
+        lexsort=LEXSORT_UNIT_COMPACT, return_order=True, return_slices=True
     )
 
-    # The buckets, in the order the requested lexsort puts them in.
-    unit_major = lexsort == ("sample_index", "segment_index", "unit_index")
-    if unit_major:
-        assert slices.shape == (num_units, num_segments, 2)
-        groups = [(u, s) for u in range(num_units) for s in range(num_segments)]
-    else:
-        assert slices.shape == (num_segments, num_units, 2)
-        groups = [(u, s) for s in range(num_segments) for u in range(num_units)]
+    # The buckets, unit by unit and segment by segment within each unit.
+    assert slices.shape == (num_units, num_segments, 2)
+    groups = [(u, s) for u in range(num_units) for s in range(num_segments)]
     masks = [(spikes["unit_index"] == u) & (spikes["segment_index"] == s) for u, s in groups]
 
     assert np.array_equal(ordered_spikes, np.concatenate([spikes[mask] for mask in masks]))

src/spikeinterface/core/tests/test_numpy_extractors.py

@@ -117,7 +117,7 @@ def test_NumpySorting(setup_NumpyRecording):
 
     # construct back from kwargs keep the same array
     sorting2 = load(sorting.to_dict())
-    assert np.shares_memory(sorting2._cached_spike_vector, sorting._cached_spike_vector)
+    assert np.shares_memory(sorting2._cached_time_ordered["spikes"], sorting._cached_time_ordered["spikes"])
 
 
 def test_SharedMemorySorting():
@@ -133,7 +133,7 @@ def test_SharedMemorySorting():
 
     sorting = SharedMemorySorting.from_sorting(np_sorting)
     # print(sorting)
-    assert sorting._cached_spike_vector is not None
+    assert sorting._cached_time_ordered is not None
 
     # print(sorting.to_spike_vector())
     d = sorting.to_dict()

src/spikeinterface/core/tests/test_unitsselectionsorting.py

@@ -147,7 +147,7 @@ def test_selection_keeps_extra_fields():
     stay with their spike through a selection."""
     wide_dtype = minimum_spike_dtype + [("channel_index", "int64")]
     parent = _make_parent_with_shuffled_ties(dtype=wide_dtype)
-    parent._cached_spike_vector["channel_index"] = np.arange(parent._cached_spike_vector.size)
+    parent._cached_time_ordered["spikes"]["channel_index"] = np.arange(parent._cached_time_ordered["spikes"].size)
 
     unit_ids = ["3", "22", "29"]
     child = UnitsSelectionSorting(parent, unit_ids=unit_ids)
@@ -188,7 +188,7 @@ def _assert_shares_parent_caches(child, parent):
     assert child._is_identity_selection is True
     assert child.to_spike_vector() is parent.to_spike_vector()
     assert child._get_spike_vector_segment_slices() is parent._get_spike_vector_segment_slices()
-    assert child._cached_lexsorted_spike_vector is parent._cached_lexsorted_spike_vector
+    assert child._cached_unit_grouped is parent._cached_unit_grouped
 
 
 @pytest.mark.parametrize("make_child", IDENTITY_SELECTIONS.values(), ids=IDENTITY_SELECTIONS.keys())
@@ -203,13 +203,13 @@ def test_identity_selection_shares_parent_cache(make_child):
     parent_reordered = parent.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT, return_order=False, return_slices=False)
     child_reordered = child.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT, return_order=False, return_slices=False)
     assert child_reordered is parent_reordered
-    assert len(parent._cached_lexsorted_spike_vector) == 1
+    assert parent._cached_unit_grouped is not None
 
     # The child computes the reordered spike vector first (via get_unit_spike_train); the parent must find it.
     parent = _make_parent_with_shuffled_ties()
     child = make_child(parent)
     child.get_unit_spike_train(child.unit_ids[0], segment_index=0)
-    assert len(parent._cached_lexsorted_spike_vector) == 1
+    assert parent._cached_unit_grouped is not None
     parent_reordered = parent.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT, return_order=False, return_slices=False)
     child_reordered = child.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT, return_order=False, return_slices=False)
     assert child_reordered is parent_reordered
@@ -241,11 +241,10 @@ def test_non_identity_selection_does_not_share(unit_ids):
     child = parent.select_units(unit_ids)
     assert child._is_identity_selection is False
     assert child.to_spike_vector() is not parent.to_spike_vector()
-    assert child._cached_lexsorted_spike_vector is not parent._cached_lexsorted_spike_vector
 
     child.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT)
-    assert len(child._cached_lexsorted_spike_vector) == 1
-    assert len(parent._cached_lexsorted_spike_vector) == 0
+    assert child._cached_unit_grouped is not None
+    assert parent._cached_unit_grouped is None
 
 
 def test_identity_selection_keeps_lazy_zarr_vector(tmp_path):

src/spikeinterface/core/unitsselectionsorting.py

@@ -34,9 +34,6 @@ def __init__(self, parent_sorting, unit_ids=None, renamed_unit_ids=None):
         BaseSorting.__init__(self, sampling_frequency, self._renamed_unit_ids)
 
         self._is_identity_selection = bool(np.array_equal(self._unit_ids, parents_unit_ids))
-        if self._is_identity_selection:
-            # Same units (possibly renamed), same order => we can use the parent's cached spike vector
-            self._cached_lexsorted_spike_vector = parent_sorting._cached_lexsorted_spike_vector
 
         for parent_segment in self._parent_sorting.segments:
             sub_segment = UnitsSelectionSortingSegment(parent_segment, ids_conversion)
@@ -50,26 +47,40 @@ def __init__(self, parent_sorting, unit_ids=None, renamed_unit_ids=None):
 
         self._kwargs = dict(parent_sorting=parent_sorting, unit_ids=unit_ids, renamed_unit_ids=renamed_unit_ids)
 
+    @property
+    def _cached_unit_grouped(self):
+        # Same units (possibly renamed), same order => we read and fill the parent's unit-grouped cache
+        if getattr(self, "_is_identity_selection", False):
+            return self._parent_sorting._cached_unit_grouped
+        return self._own_cached_unit_grouped
+
+    @_cached_unit_grouped.setter
+    def _cached_unit_grouped(self, value):
+        if getattr(self, "_is_identity_selection", False):
+            self._parent_sorting._cached_unit_grouped = value
+        else:
+            self._own_cached_unit_grouped = value
+
     def _compute_and_cache_spike_vector(self) -> None:
         from spikeinterface.core.sorting_tools import remap_unit_indices_in_vector
 
-        if self._parent_sorting._cached_spike_vector is None:
+        if self._parent_sorting._cached_time_ordered is None:
             self._parent_sorting._compute_and_cache_spike_vector()
 
-            if self._parent_sorting._cached_spike_vector is None:
+            if self._parent_sorting._cached_time_ordered is None:
                 return
 
         if self._is_identity_selection:
-            self._cached_spike_vector = self._parent_sorting._cached_spike_vector
-            self._cached_spike_vector_segment_slices = self._parent_sorting._get_spike_vector_segment_slices()
+            self._parent_sorting._get_spike_vector_segment_slices()
+            self._cached_time_ordered = self._parent_sorting._cached_time_ordered
             return
 
         spike_vector, _ = remap_unit_indices_in_vector(
-            vector=self._parent_sorting._cached_spike_vector,
+            vector=self._parent_sorting._cached_time_ordered["spikes"],
             all_old_unit_ids=self._parent_sorting.unit_ids,
             all_new_unit_ids=self._unit_ids,
         )
-        self._cached_spike_vector = spike_vector
+        self._cached_time_ordered = {"spikes": spike_vector, "segment_slices": None}
 
     def search_cached_spikes_sorted(
         self,

src/spikeinterface/core/zarrextractors.py

@@ -608,11 +608,12 @@ def __init__(
         self._lazy_spike_vector = lazy_spike_vector
         self._spikes_group = spikes_group
         self._sample_index_search = None
-
-        self._cached_spike_vector = spikes
         # pre-populate segment slices so _get_spike_vector_segment_slices() never
         # needs to materialise the full segment_index array
-        self._cached_spike_vector_segment_slices = segment_slices_list
+        self._cached_time_ordered = {
+            "spikes": spikes,
+            "segment_slices": np.asarray(segment_slices_list, dtype="int64"),
+        }
 
         for segment_index in range(num_segments):
             soring_segment = SpikeVectorSortingSegment(spikes, segment_index, unit_ids)
@@ -660,7 +661,7 @@ def search_cached_spikes_sorted(
             if "sample_index_chunk_firsts" in self._spikes_group:
                 chunk_firsts = self._spikes_group["sample_index_chunk_firsts"][:]
             self._sample_index_search = ZarrSampleIndexSearch(self._spikes_group["sample_index"], chunk_firsts)
-        start, stop = self._cached_spike_vector_segment_slices[segment_index]
+        start, stop = self._cached_time_ordered["segment_slices"][segment_index]
         return self._sample_index_search.searchsorted(indices, int(start), int(stop))
 
     @staticmethod
@@ -896,10 +897,10 @@ def add_sorting_to_zarr_group(
             codec_kwargs = build_codec_pipeline(filters=filters, compressors=compressors)
             spikes_group.create_array(name=field, data=spikes[field], chunks=chunks, shards=shards, **codec_kwargs)
 
-    segment_slices = []
-    if sorting._cached_spike_vector_segment_slices is not None:
-        segment_slices = sorting._cached_spike_vector_segment_slices
+    if sorting._cached_time_ordered.get("segment_slices") is not None:
+        segment_slices = sorting._cached_time_ordered["segment_slices"]
     else:
+        segment_slices = []
         for segment_index in range(num_segments):
             i0, i1 = np.searchsorted(spikes["segment_index"], [segment_index, segment_index + 1])
             segment_slices.append([i0, i1])

src/spikeinterface/curation/remove_excess_spikes.py

@@ -46,13 +46,13 @@ def __init__(self, sorting: BaseSorting, recording: BaseRecording) -> None:
         self._kwargs = {"sorting": sorting, "recording": recording}
 
     def _compute_and_cache_spike_vector(self) -> None:
-        if self._parent_sorting._cached_spike_vector is None:
+        if self._parent_sorting._cached_time_ordered is None:
             self._parent_sorting._compute_and_cache_spike_vector()
 
-            if self._parent_sorting._cached_spike_vector is None:
+            if self._parent_sorting._cached_time_ordered is None:
                 return
 
-        parent_spike_vector = self._parent_sorting._cached_spike_vector
+        parent_spike_vector = self._parent_sorting._cached_time_ordered["spikes"]
         num_segments = self._parent_sorting.get_num_segments()
 
         list_spike_vectors = []
@@ -64,7 +64,7 @@ def _compute_and_cache_spike_vector(self) -> None:
             list_spike_vectors.append(spike_vector[start:end])
 
         spike_vector = np.concatenate(list_spike_vectors)
-        self._cached_spike_vector = spike_vector
+        self._cached_time_ordered = {"spikes": spike_vector, "segment_slices": None}
 
 
 class RemoveExcessSpikesSortingSegment(BaseSortingSegment):

src/spikeinterface/extractors/nwbextractors.py

@@ -1550,8 +1550,10 @@ def _compute_and_cache_spike_vector(self) -> None:
         spikes["sample_index"] = sample_indices[order]
         spikes["unit_index"] = unit_indices[order]
         # segment_index stays 0 for the single segment.
-        self._cached_spike_vector = spikes
-        self._cached_spike_vector_segment_slices = np.array([[0, num_spikes]], dtype="int64")
+        self._cached_time_ordered = {
+            "spikes": spikes,
+            "segment_slices": np.array([[0, num_spikes]], dtype="int64"),
+        }
 
     def get_last_spike_frame(self, segment_index: int | None = None) -> int:
         # Performance: the generic BaseSorting version builds the full spike vector to take one maximum.

src/spikeinterface/extractors/phykilosortextractors.py

@@ -268,9 +268,10 @@ def _compute_and_cache_spike_vector(self) -> None:
         # This is useless because phy is always one segment
         # spikes["segment_index"] = 0
 
-        self._cached_spike_vector = spikes
-        self._cached_spike_vector_segment_slices = np.zeros((1, 2), dtype="int64")
-        self._cached_spike_vector_segment_slices[0, 1] = n
+        self._cached_time_ordered = {
+            "spikes": spikes,
+            "segment_slices": np.array([[0, n]], dtype="int64"),
+        }
 
 
 class PhySortingSegment(BaseSortingSegment):