diff --git a/src/spikeinterface/core/tests/test_unitsselectionsorting.py b/src/spikeinterface/core/tests/test_unitsselectionsorting.py index 0b9cf6c813..309563df62 100644 --- a/src/spikeinterface/core/tests/test_unitsselectionsorting.py +++ b/src/spikeinterface/core/tests/test_unitsselectionsorting.py @@ -248,6 +248,7 @@ def test_non_identity_selection_does_not_share(unit_ids): assert len(parent._cached_lexsorted_spike_vector) == 0 +@pytest.mark.requires_zarr_write def test_identity_selection_keeps_lazy_zarr_vector(tmp_path): """A lazy parent spike vector should stay lazy through an identity selection.""" from spikeinterface.core.zarrextractors import ZarrSortingExtractor, ZarrSpikeVector @@ -262,6 +263,11 @@ def test_identity_selection_keeps_lazy_zarr_vector(tmp_path): assert renamed._get_spike_vector_segment_slices() is lazy_parent._get_spike_vector_segment_slices() eager_parent = ZarrSortingExtractor(folder) + frames = [-1, 0, 5, 7, 10_000] + np.testing.assert_array_equal( + renamed.search_cached_spikes_sorted(frames, segment_index=0), + eager_parent.search_cached_spikes_sorted(frames, segment_index=0), + ) subset = lazy_parent.select_units(["29", "22"]) assert isinstance(subset.to_spike_vector(), np.ndarray) assert np.array_equal(subset.to_spike_vector(), _mask_and_remap(eager_parent, ["29", "22"])) diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index e1b315e592..2b388839d1 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -102,6 +102,7 @@ def test_ZarrSampleIndexSearch(tmp_path): np.testing.assert_array_equal(search.searchsorted([5], 10, 10), [0]) +@pytest.mark.requires_zarr_write def test_ZarrSortingExtractor_lazy_search(tmp_path): sorting = generate_sorting(num_units=10, durations=[5.0, 3.0, 4.0], firing_rates=40.0, seed=0) folder = tmp_path / "sorting.zarr" diff --git a/src/spikeinterface/core/unitsselectionsorting.py b/src/spikeinterface/core/unitsselectionsorting.py index 31fbe2098f..33fec6fa94 100644 --- a/src/spikeinterface/core/unitsselectionsorting.py +++ b/src/spikeinterface/core/unitsselectionsorting.py @@ -71,6 +71,16 @@ def _compute_and_cache_spike_vector(self) -> None: ) self._cached_spike_vector = spike_vector + def search_cached_spikes_sorted( + self, + indices: list[int], + segment_index: int | None = None, + ): + if self._is_identity_selection: + # same spikes in the same positions, and the parent knows how to search a lazy vector + return self._parent_sorting.search_cached_spikes_sorted(indices=indices, segment_index=segment_index) + return super().search_cached_spikes_sorted(indices=indices, segment_index=segment_index) + class UnitsSelectionSortingSegment(BaseSortingSegment): def __init__(self, parent_segment, ids_conversion):