diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py index 261710278a..dde1ba8620 100644 --- a/src/spikeinterface/core/analyzer_extension_core.py +++ b/src/spikeinterface/core/analyzer_extension_core.py @@ -1618,7 +1618,7 @@ def _get_data(self, outputs="numpy", concatenated=False, return_data_name=None, sorting = self.sorting_analyzer.sorting if outputs == "numpy": - if copy: + if copy and not self.sorting_analyzer._lazy: return all_data.copy() # return a copy to avoid modification else: return all_data diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index a7af4d9630..62d5369b0d 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -357,7 +357,9 @@ def create_sorting_analyzer( return sorting_analyzer -def load_sorting_analyzer(folder, load_extensions=True, format="auto", backend_options=None) -> "SortingAnalyzer": +def load_sorting_analyzer( + folder, load_extensions=True, format="auto", backend_options=None, lazy=False +) -> "SortingAnalyzer": """ Load a SortingAnalyzer object from disk. @@ -385,7 +387,9 @@ def load_sorting_analyzer(folder, load_extensions=True, format="auto", backend_o The loaded SortingAnalyzer """ - return SortingAnalyzer.load(folder, load_extensions=load_extensions, format=format, backend_options=backend_options) + return SortingAnalyzer.load( + folder, load_extensions=load_extensions, format=format, backend_options=backend_options, lazy=lazy + ) class SortingAnalyzer: @@ -421,6 +425,7 @@ def __init__( peak_sign: PeakSignType = "both", peak_mode: PeakModeType = "extremum", backend_options: dict | None = None, + lazy: bool = False, ): # very fast init because checks are done in load and create self.sorting = sorting @@ -449,6 +454,9 @@ def __init__( # (additional saving options for creating and saving datasets, e.g. compression/filters for zarr) self._backend_options = {} if backend_options is None else backend_options + # the lazy flag is used to load the extensions in a lazy way (only when needed) + self._lazy = lazy + # extensions are not loaded at init self.extensions = dict() @@ -581,6 +589,7 @@ def load( load_extensions: bool = True, format: Literal["auto", "binary_folder", "zarr"] = "auto", backend_options: dict | None = None, + lazy: bool = False, ): """ Load folder or zarr. @@ -594,16 +603,16 @@ def load( if format == "binary_folder": sorting_analyzer = SortingAnalyzer.load_from_binary_folder( - folder, recording=recording, backend_options=backend_options + folder, recording=recording, backend_options=backend_options, lazy=lazy ) elif format == "zarr": sorting_analyzer = SortingAnalyzer.load_from_zarr( - folder, recording=recording, backend_options=backend_options + folder, recording=recording, backend_options=backend_options, lazy=lazy ) else: raise ValueError(f"SortingAnalyzer.load: wrong format {format}") - if load_extensions and not is_path_remote(folder): + if load_extensions and not lazy and not is_path_remote(folder): sorting_analyzer.load_all_saved_extension() return sorting_analyzer @@ -885,6 +894,7 @@ def load_from_binary_folder( folder: str | Path, recording: BaseRecording | None = None, backend_options: dict | None = None, + lazy: bool = False, ) -> "SortingAnalyzer": from .loading import load @@ -928,9 +938,18 @@ def load_from_binary_folder( with open(settings_file, "w") as f: json.dump(check_json(settings), f, indent=4) - # Load sorting (in memory) + # Load sorting (in memory or lazy) + if lazy: + numpy_folder_kwargs = dict(mmap_mode="r") + copy_spike_vector = False + else: + numpy_folder_kwargs = dict() + copy_spike_vector = True + sorting = NumpySorting.from_sorting( - NumpyFolderSorting(sorting_folder), with_metadata=True, copy_spike_vector=True + NumpyFolderSorting(folder / "sorting", **numpy_folder_kwargs), + with_metadata=True, + copy_spike_vector=copy_spike_vector, ) # Load recording (if available) @@ -970,6 +989,7 @@ def load_from_binary_folder( peak_sign=settings["peak_sign"], peak_mode=settings["peak_mode"], backend_options=backend_options, + lazy=lazy, ) sorting_analyzer.folder = folder @@ -1088,6 +1108,7 @@ def load_from_zarr( folder: str | Path, recording: BaseRecording | None = None, backend_options: dict | None = None, + lazy: bool = False, ) -> "SortingAnalyzer": import zarr from .loading import load @@ -1117,11 +1138,22 @@ def load_from_zarr( settings = zarr_root.attrs["settings"] settings = cls._handle_backward_compatibility_settings_pre_init(settings) - # Load sorting (in memory) + # Load sorting (in memory or lazy) + if lazy: + copy_spike_vector = False + lazy_spike_vector = True + else: + copy_spike_vector = True + lazy_spike_vector = False sorting = NumpySorting.from_sorting( - ZarrSortingExtractor(folder, zarr_group="sorting", storage_options=storage_options), + ZarrSortingExtractor( + folder, + zarr_group="sorting", + storage_options=storage_options, + lazy_spike_vector=lazy_spike_vector, + ), with_metadata=True, - copy_spike_vector=True, + copy_spike_vector=copy_spike_vector, ) # Load recording (if available) @@ -1161,6 +1193,7 @@ def load_from_zarr( peak_sign=settings["peak_sign"], peak_mode=settings["peak_mode"], backend_options=backend_options, + lazy=lazy, ) sorting_analyzer.folder = folder @@ -1486,6 +1519,11 @@ def _save_or_select_or_merge_or_split( new_sorting_analyzer : SortingAnalyzer The newly created SortingAnalyzer object. """ + if self._lazy: + raise ValueError( + "Cannot save, select, merge or split units when the SortingAnalyzer is lazy. " + "Please load the SortingAnalyzer with lazy=False." + ) if self.has_recording(): recording = self._recording elif self.has_temporary_recording(): @@ -2212,6 +2250,10 @@ def compute(self, input, save=True, extension_params=None, verbose=False, **kwar ) """ + if self._lazy: + # If the analyzer is lazy, we can compute extensions in memory but we won't save / overwrite any existing + # extension on disk. This is to avoid overwriting existing extensions when the analyzer is lazy. + save = False if isinstance(input, str): return self.compute_one_extension(extension_name=input, save=save, verbose=verbose, **kwargs) elif isinstance(input, dict): @@ -2508,7 +2550,7 @@ def load_extension(self, extension_name: str): if extension_class is None: return None - extension_instance = extension_class.load(self) + extension_instance = extension_class.load(self, lazy=self._lazy) self.extensions[extension_name] = extension_instance @@ -2527,7 +2569,7 @@ def delete_extension(self, extension_name) -> None: """ # delete from folder or zarr - if self.format != "memory" and self.has_extension(extension_name): + if self.format != "memory" and self.has_extension(extension_name) and not self._lazy: # need a reload to reset the folder ext = self.load_extension(extension_name) ext.delete() @@ -2990,20 +3032,20 @@ def _get_zarr_extension_group(self, mode="r+"): return extension_group @classmethod - def load(cls, sorting_analyzer): + def load(cls, sorting_analyzer, lazy=False): ext = cls(sorting_analyzer) ext.load_params() ext.load_run_info() if ext.run_info is not None: if ext.run_info["run_completed"]: - ext.load_data() + ext.load_data(lazy=lazy) if cls.need_backward_compatibility_on_load: ext._handle_backward_compatibility_on_load() if len(ext.data) > 0: return ext else: # this is for back-compatibility of old analyzers - ext.load_data() + ext.load_data(lazy=lazy) if cls.need_backward_compatibility_on_load: ext._handle_backward_compatibility_on_load() if len(ext.data) > 0: @@ -3103,7 +3145,7 @@ def load_params(self): self.params = params - def load_data(self): + def load_data(self, lazy=False): ext_data = None if self.format == "binary_folder": extension_folder = self._get_binary_extension_folder() @@ -3123,10 +3165,12 @@ def load_data(self): ext_data = json.load(f) elif ext_data_file.suffix == ".npy": # The lazy loading of an extension is complicated because if we compute again - # and have a link to the old buffer on windows then it fails - # ext_data = np.load(ext_data_file, mmap_mode="r") - # so we go back to full loading - ext_data = np.load(ext_data_file) + # and have a link to the old buffer on windows then it fails. + # So, by default, we use full loading, but lazy can be requested on demand. + if lazy: + ext_data = np.load(ext_data_file, mmap_mode="r") + else: + ext_data = np.load(ext_data_file) elif ext_data_file.suffix == ".csv": import pandas as pd @@ -3162,8 +3206,7 @@ def load_data(self): elif "object" in ext_data_.attrs: ext_data = ext_data_[0] else: - # this load in memory - ext_data = np.array(ext_data_) + ext_data = ext_data_ if lazy else np.array(ext_data_[:]) self.set_data(ext_data_name, ext_data) if len(self.data) == 0: diff --git a/src/spikeinterface/core/sortingfolder.py b/src/spikeinterface/core/sortingfolder.py index 6fc4729bf2..ad331cf835 100644 --- a/src/spikeinterface/core/sortingfolder.py +++ b/src/spikeinterface/core/sortingfolder.py @@ -24,7 +24,7 @@ class NumpyFolderSorting(BaseSorting): mode = "folder" name = "NumpyFolder" - def __init__(self, folder_path: str | Path): + def __init__(self, folder_path, mmap_mode: str | None = None): folder_path = Path(folder_path) # Load general info @@ -37,8 +37,8 @@ def __init__(self, folder_path: str | Path): # Init superclass super().__init__(sampling_frequency, unit_ids) - # Load spikes vector - self.spikes = np.load(folder_path / "spikes.npy") + self.spikes = np.load(folder_path / "spikes.npy", mmap_mode=mmap_mode) + 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 @@ -47,8 +47,7 @@ def __init__(self, folder_path: str | Path): # Load metadata self.load_metadata_from_folder(folder_path) - # Save folder_path as kwargs for serialization - self._kwargs = {"folder_path": str(folder_path.absolute())} + self._kwargs = dict(folder_path=str(folder_path.absolute()), mmap_mode=mmap_mode) @staticmethod def write_sorting(sorting, save_path): diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index f293699e3b..25aeb78c1a 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -131,7 +131,7 @@ def test_SortingAnalyzer_binary_folder(tmp_path, dataset): assert "number" in sorting_analyzer.sorting.get_property_keys() sorting_analyzer_reloded = load_sorting_analyzer(folder, format="auto") assert "quality" in sorting_analyzer_reloded.sorting.get_property_keys() - assert "number" in sorting_analyzer.sorting.get_property_keys() + assert "number" in sorting_analyzer_reloded.sorting.get_property_keys() def test_SortingAnalyzer_zarr(tmp_path, dataset): @@ -213,7 +213,7 @@ def test_SortingAnalyzer_zarr(tmp_path, dataset): assert "number" in sorting_analyzer.sorting.get_property_keys() sorting_analyzer_reloded = load_sorting_analyzer(sorting_analyzer.folder, format="auto") assert "quality" in sorting_analyzer_reloded.sorting.get_property_keys() - assert "number" in sorting_analyzer.sorting.get_property_keys() + assert "number" in sorting_analyzer_reloded.sorting.get_property_keys() def test_create_by_dict(): @@ -361,6 +361,53 @@ def test_SortingAnalyzer_interleaved_probegroup(dataset): assert np.array_equal(recording.get_channel_locations(), sorting_analyzer.get_channel_locations()) +@pytest.mark.parametrize("format", ["binary_folder", "zarr"]) +def test_load_in_lazy_mode(tmp_path, dataset, format): + recording, sorting = dataset + + folder = tmp_path / "test_SortingAnalyzer_folder" + if format == "zarr": + import zarr + from spikeinterface.core.zarrextractors import ZarrSpikeVector + + folder = folder.with_suffix(".zarr") + array_class = zarr.Array + spike_vector_class = ZarrSpikeVector + else: + array_class = np.memmap + spike_vector_class = np.memmap + if folder.exists(): + shutil.rmtree(folder) + + sorting_analyzer = create_sorting_analyzer( + sorting, recording, format=format, folder=folder, sparse=False, sparsity=None + ) + + sorting_analyzer.compute(["random_spikes", "templates", "spike_amplitudes"]) + # load in lazy mode and check that spike vector and extension data are memmap + sorting_analyzer_lazy = load_sorting_analyzer(folder, format="auto", lazy=True) + + assert isinstance(sorting_analyzer_lazy.sorting.to_spike_vector(), spike_vector_class) + + template_ext = sorting_analyzer_lazy.get_extension("templates") + template_data = template_ext.data + for key, value in template_data.items(): + if isinstance(value, np.ndarray): + assert isinstance(value, array_class) + spike_amplitudes_ext = sorting_analyzer_lazy.get_extension("spike_amplitudes") + spike_amplitudes_data = spike_amplitudes_ext.data + for key, value in spike_amplitudes_data.items(): + if isinstance(value, np.ndarray): + assert isinstance(value, array_class) + + # check that the lazy mode does not overwrite existing extensions + sorting_analyzer_lazy.compute("random_spikes", max_spikes_per_unit=10) + # reload the analyzer to check that the original extension is not overwritten + sorting_analyzer_reloaded = load_sorting_analyzer(folder, format="auto", lazy=True) + random_spikes_ext = sorting_analyzer_reloaded.get_extension("random_spikes") + assert random_spikes_ext.params["max_spikes_per_unit"] != 10 + + def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder): register_result_extension(DummyAnalyzerExtension) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index b41ae26827..02ae5e9ecd 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -268,6 +268,112 @@ def get_traces( return traces +class _ZarrSegmentIndex: + """Lazy segment_index array derived from segment_slices stored in zarr.""" + + def __init__(self, segment_slices: np.ndarray, n: int): + self._segment_slices = segment_slices + self._n = n + + def __len__(self) -> int: + return self._n + + def __array__(self, dtype=None): + arr = np.empty(self._n, dtype="int64") + for seg_idx, (s0, s1) in enumerate(self._segment_slices): + arr[s0:s1] = seg_idx + return arr if dtype is None else arr.astype(dtype) + + def __getitem__(self, key): + return np.asarray(self)[key] + + def __eq__(self, other): + return np.asarray(self) == other + + +class ZarrSpikeVector: + """ + Virtual structured spike vector backed by zarr arrays. + + Mimics a memmap-backed numpy structured array with fields + (sample_index, unit_index, segment_index) without loading any data + at construction time. Data is read from zarr lazily: + + * Field access (``spikes["sample_index"]``) returns the zarr array + (or a lazy segment-index object). + * Slice access (``spikes[s0:s1]``) materialises only that slice. + * ``np.asarray(spikes)`` materialises the full array. + + The zarr arrays are assumed to be stored in sorted order + (segment_index ASC, sample_index ASC, unit_index ASC), which is the + ordering guaranteed by :func:`add_sorting_to_zarr_group`. + """ + + def __init__(self, spikes_group, segment_slices: np.ndarray): + self._sample_index = spikes_group["sample_index"] + self._unit_index = spikes_group["unit_index"] + self._segment_slices = np.asarray(segment_slices, dtype="int64") + self._n = len(self._sample_index) + self.dtype = np.dtype(minimum_spike_dtype) + + @property + def size(self) -> int: + return self._n + + def __len__(self) -> int: + return self._n + + def __getitem__(self, key): + if isinstance(key, str): + if key == "sample_index": + return self._sample_index + elif key == "unit_index": + return self._unit_index + elif key == "segment_index": + return _ZarrSegmentIndex(self._segment_slices, self._n) + else: + raise KeyError(f"ZarrSpikeVector has no field {key!r}") + + if isinstance(key, (int, np.integer)): + idx = int(key) + if idx < 0: + idx += self._n + result = np.empty(1, dtype=self.dtype) + result["sample_index"][0] = self._sample_index[idx] + result["unit_index"][0] = self._unit_index[idx] + result["segment_index"][0] = int(np.searchsorted(self._segment_slices[:, 0], idx, side="right")) - 1 + return result[0] + + if isinstance(key, slice): + start, stop, step = key.indices(self._n) + n = len(range(start, stop, step)) + result = np.empty(n, dtype=self.dtype) + result["sample_index"] = self._sample_index[start:stop:step] + result["unit_index"] = self._unit_index[start:stop:step] + if step == 1: + seg_index = np.empty(n, dtype="int64") + for seg_idx, (s0, s1) in enumerate(self._segment_slices): + lo = max(start, int(s0)) - start + hi = min(stop, int(s1)) - start + if hi > lo: + seg_index[lo:hi] = seg_idx + result["segment_index"] = seg_index + else: + result["segment_index"] = _ZarrSegmentIndex(self._segment_slices, self._n)[start:stop:step] + return result + + # fallback for fancy/boolean indexing: materialise then index + return np.asarray(self)[key] + + def __array__(self, dtype=None): + arr = np.empty(self._n, dtype=self.dtype) + arr["sample_index"] = self._sample_index[:] + arr["unit_index"] = self._unit_index[:] + for seg_idx, (s0, s1) in enumerate(self._segment_slices): + arr["segment_index"][s0:s1] = seg_idx + return arr if dtype is None else arr.astype(dtype) + + class ZarrSortingExtractor(BaseSorting): """ SortingExtractor for a zarr format @@ -284,13 +390,23 @@ class ZarrSortingExtractor(BaseSorting): Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc. zarr_group : str or None, default: None Optional zarr group path to load the sorting from. This can be used when the sorting is not stored at the root, but in sub group. + lazy_spike_vector : bool, default: False + If True, the spike vector is loaded lazily. This can be useful for large sortings with many spikes. + If False, the spike vector is loaded in memory. Default: False + Returns ------- sorting : ZarrSortingExtractor The sorting Extractor """ - def __init__(self, folder_path: Path | str, storage_options: dict | None = None, zarr_group: str | None = None): + def __init__( + self, + folder_path: Path | str, + storage_options: dict | None = None, + zarr_group: str | None = None, + lazy_spike_vector: bool = False, + ): folder_path, folder_path_kwarg = resolve_zarr_path(folder_path) @@ -316,16 +432,23 @@ def __init__(self, folder_path: Path | str, storage_options: dict | None = None, BaseSorting.__init__(self, sampling_frequency, unit_ids) - spikes = np.zeros(len(spikes_group["sample_index"]), dtype=minimum_spike_dtype) - spikes["sample_index"] = spikes_group["sample_index"][:] - spikes["unit_index"] = spikes_group["unit_index"][:] - for i, (start, end) in enumerate(segment_slices_list): - spikes["segment_index"][start:end] = i - # we do not need to lexsort at init (very high cost) because there already sorted by frame before to be saved. - # In version 0.104.X this was fully lexsorted, but we don't need it anymore because it's only important in the context of SpikeVectorBased extensions in the SortingAnalyzer, which stores its own copy of the Sorting object. This makes the extension data and the spike vector always matching their order. - # spikes = spikes[np.lexsort((spikes["unit_index"], spikes["sample_index"], spikes["segment_index"]))] + if lazy_spike_vector: + spikes = ZarrSpikeVector(spikes_group, segment_slices_list) + else: + # Materialize the spike vector in memory and sort it by (segment_index, sample_index, unit_index) + spikes = np.zeros(len(spikes_group["sample_index"]), dtype=minimum_spike_dtype) + spikes["sample_index"] = spikes_group["sample_index"][:] + spikes["unit_index"] = spikes_group["unit_index"][:] + for i, (start, end) in enumerate(segment_slices_list): + spikes["segment_index"][start:end] = i + # we do not need to lexsort at init (very high cost) because there already sorted by frame before to be saved. + # In version 0.104.X this was fully lexsorted, but we don't need it anymore because it's only important in the context of SpikeVectorBased extensions in the SortingAnalyzer, which stores its own copy of the Sorting object. This makes the extension data and the spike vector always matching their order. + # spikes = spikes[np.lexsort((spikes["unit_index"], spikes["sample_index"], spikes["segment_index"]))] 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 = np.asarray(segment_slices_list, dtype="int64") for segment_index in range(num_segments): soring_segment = SpikeVectorSortingSegment(spikes, segment_index, unit_ids) @@ -343,7 +466,12 @@ def __init__(self, folder_path: Path | str, storage_options: dict | None = None, if annotations is not None: self.annotate(**annotations) - self._kwargs = {"folder_path": folder_path_kwarg, "storage_options": storage_options, "zarr_group": zarr_group} + self._kwargs = { + "folder_path": folder_path_kwarg, + "storage_options": storage_options, + "zarr_group": zarr_group, + "lazy_spike_vector": lazy_spike_vector, + } @staticmethod def write_sorting(sorting: BaseSorting, folder_path: str | Path, storage_options: dict | None = None, **kwargs):