diff --git a/doc/modules/core.rst b/doc/modules/core.rst index 6ea3d25eb6..db6c2534c4 100644 --- a/doc/modules/core.rst +++ b/doc/modules/core.rst @@ -274,8 +274,24 @@ Once a :code:`SortingAnalyzer` object is saved to disk, it can be easily reloade .. code-block:: python + sorting_analyzer = si.load(folder="my-sorting-analyzer.zarr") + # or sorting_analyzer = si.load_sorting_analyzer(folder="my-sorting-analyzer.zarr") +The :code:`SortingAnalyzer` can also be loaded in :code:`read_only` mode, which prevents any modification of the +:code:`SortingAnalyzer` or its extensions. + +.. code-block:: python + + sorting_analyzer = si.load_sorting_analyzer(folder="my-sorting-analyzer.zarr", read_only=True) + +Finally, the :code:`SortingAnalyzer` can be loaded in :code:`lazy` mode, which will not load the extensions data in +memory as arrays, but keep them as memmap or zarr arrays. This is useful for very large datasets, where the extensions +data can be very large and not fit in memory, but it can be slower to access the data. + +.. code-block:: python + + sorting_analyzer = si.load_sorting_analyzer(folder="my-sorting-analyzer.zarr", lazy=True) .. note:: @@ -420,6 +436,24 @@ All computed extensions will be automatically propagated or merged when curating :ref:`modules/curation:Curation module` documentation for more information. +Handling very large datasets: ``lazy`` mode +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +For very large datasets with tens-to-hundreds millions of spikes, the :code:`SortingAnalyzer` computations can be very +memory intensive. By default, in fact, the :code:`SortingAnalyzer` computes and stores all the data in memory. +However, it is possible to use a :code:`lazy` mode, which will compute the data on-the-fly and store them on-disk as they +are computed. This makes some computations slower, but it allows to handle very large datasets without running out of memory. +Note that the :code:`lazy` mode is only available for the :code:`zarr` and :code:`binary_folder` backends. + +.. code-block:: python + + sorting_analyzer_lazy = create_sorting_analyzer( + sorting=sorting, + recording=recording, + format="zarr", + lazy=True, # compute on-the-fly and store on-disk + ) + Event ----- diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py index dde1ba8620..7fee9a5bc6 100644 --- a/src/spikeinterface/core/analyzer_extension_core.py +++ b/src/spikeinterface/core/analyzer_extension_core.py @@ -21,7 +21,7 @@ from .template import Templates from .sorting_tools import random_spikes_selection, select_sorting_periods_mask, spike_vector_to_indices from .job_tools import fix_job_kwargs, split_job_kwargs -from .core_tools import ms_to_samples +from .core_tools import ms_to_samples, slice_rows, materialize_array class ComputeRandomSpikes(AnalyzerExtension): @@ -96,7 +96,8 @@ def _merge_extension_data( new_data = dict() random_spikes_indices = self.data["random_spikes_indices"] if keep_mask is None: - new_data["random_spikes_indices"] = random_spikes_indices.copy() + # no filtering: sharing the reference is fine, materialization happens on save + new_data["random_spikes_indices"] = random_spikes_indices else: spikes = self.sorting_analyzer.sorting.to_spike_vector() selected_mask = np.zeros(spikes.size, dtype=bool) @@ -106,7 +107,8 @@ def _merge_extension_data( def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs): new_data = dict() - new_data["random_spikes_indices"] = self.data["random_spikes_indices"].copy() + # no filtering: sharing the reference is fine, materialization happens on save + new_data["random_spikes_indices"] = self.data["random_spikes_indices"] return new_data def _get_data(self): @@ -253,7 +255,7 @@ def _select_units_extension_data(self, unit_ids): keep_spike_mask = np.isin(some_spikes["unit_index"], keep_unit_indices) new_data = dict() - new_data["waveforms"] = self.data["waveforms"][keep_spike_mask, :, :] + new_data["waveforms"] = slice_rows(self.data["waveforms"], keep_spike_mask) return new_data @@ -266,12 +268,15 @@ def _merge_extension_data( spike_indices = self.sorting_analyzer.get_extension("random_spikes").get_data() valid = keep_mask[spike_indices] some_spikes = some_spikes[valid] - waveforms = waveforms[valid] - else: - waveforms = waveforms.copy() + # slice_rows already returns an independent, materialized array + waveforms = slice_rows(waveforms, valid) old_sparsity = self.sorting_analyzer.sparsity if old_sparsity is not None: + if keep_mask is None: + # about to mutate waveforms in place below (sparse realignment): we need a genuinely + # independent, writable buffer rather than sharing the original reference/zarr handle + waveforms = materialize_array(waveforms) # we need a realignement inside each group because we take the channel intersection sparsity for group_ids in merge_unit_groups: group_indices = self.sorting_analyzer.sorting.ids_to_indices(group_ids) @@ -291,8 +296,9 @@ def _merge_extension_data( return dict(waveforms=waveforms) def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs): - # splitting only affects random spikes, not waveforms - new_data = dict(waveforms=self.data["waveforms"].copy()) + # splitting only affects random spikes, not waveforms: sharing the reference is fine, + # materialization happens on save + new_data = dict(waveforms=self.data["waveforms"]) return new_data def get_waveforms_one_unit(self, unit_id, force_dense: bool = False): @@ -319,7 +325,7 @@ def get_waveforms_one_unit(self, unit_id, force_dense: bool = False): some_spikes = self.sorting_analyzer.get_extension("random_spikes").get_random_spikes() spike_mask = some_spikes["unit_index"] == unit_index - wfs = waveforms[spike_mask, :, :] + wfs = slice_rows(waveforms, spike_mask) if self.sorting_analyzer.sparsity is not None: chan_inds = self.sorting_analyzer.sparsity.unit_id_to_channel_indices[unit_id] @@ -557,7 +563,7 @@ def _select_units_extension_data(self, unit_ids): new_data = dict() for key, arr in self.data.items(): - new_data[key] = arr[keep_unit_indices, :, :] + new_data[key] = slice_rows(arr, keep_unit_indices) return new_data @@ -566,7 +572,7 @@ def _select_channels_extension_data(self, channel_ids): new_data = {} for key, arr in self.data.items(): - new_data[key] = arr[:, :, keep_channel_indices] + new_data[key] = slice_rows(arr, keep_channel_indices, axis=2) return new_data @@ -591,9 +597,9 @@ def _merge_extension_data( for count, merge_unit_id in enumerate(merge_group): weights[count] = counts[merge_unit_id] weights /= weights.sum() - new_data[key][unit_index] = (arr[keep_unit_indices, :, :] * weights[:, np.newaxis, np.newaxis]).sum( - 0 - ) + new_data[key][unit_index] = ( + slice_rows(arr, keep_unit_indices) * weights[:, np.newaxis, np.newaxis] + ).sum(0) if new_sorting_analyzer.sparsity is not None: chan_ids = new_sorting_analyzer.sparsity.unit_id_to_channel_indices[unit_id] mask = ~np.isin(np.arange(arr.shape[2]), chan_ids) @@ -616,7 +622,7 @@ def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, unsplit_unit_ids = [unit_id for unit_id in self.sorting_analyzer.unit_ids if unit_id not in split_units] new_indices = np.array([new_analyzer_unit_ids.index(unit_id) for unit_id in unsplit_unit_ids]) old_indices = self.sorting_analyzer.sorting.ids_to_indices(unsplit_unit_ids) - new_array[new_indices, ...] = arr[old_indices, ...] + new_array[new_indices, ...] = slice_rows(arr, old_indices) for split_unit_id, new_splits in zip(split_units, new_unit_ids): if new_sorting_analyzer.has_extension("waveforms"): @@ -800,6 +806,13 @@ class ComputeNoiseLevels(AnalyzerExtension): def _set_params(self, **noise_level_params): params = noise_level_params.copy() + # ensure that random_slices_kwargs is always present and has a seed for reproducibility + if "random_slices_kwargs" not in params: + params["random_slices_kwargs"] = dict() + if params["random_slices_kwargs"].get("seed") is None: + from spikeinterface.core.core_tools import _ensure_seed + + params["random_slices_kwargs"]["seed"] = _ensure_seed(params["random_slices_kwargs"].get("seed")) return params def _select_units_extension_data(self, unit_ids): @@ -809,7 +822,7 @@ def _select_units_extension_data(self, unit_ids): def _select_channels_extension_data(self, channel_ids): # this does not depend on channels channel_indices = self.sorting_analyzer.channel_ids_to_indices(channel_ids) - return dict(noise_levels=self.data["noise_levels"][channel_indices]) + return dict(noise_levels=slice_rows(self.data["noise_levels"], channel_indices)) def _merge_extension_data( self, merge_unit_groups, new_unit_ids, new_sorting_analyzer, keep_mask=None, verbose=False, **job_kwargs @@ -1215,7 +1228,7 @@ def _set_params( ) return params - def _prepare_data(self, sorting_analyzer, unit_ids=None): + def _prepare_data(self, sorting_analyzer, unit_ids=None, periods=None): """ Optional function to prepare shared data for metric computation. @@ -1257,11 +1270,11 @@ def _compute_metrics( if unit_ids is None: unit_ids = sorting_analyzer.unit_ids - tmp_data = self._prepare_data(sorting_analyzer=sorting_analyzer, unit_ids=unit_ids) if metric_names is None: metric_names = self.params["metric_names"] periods = self.params.get("periods", None) + tmp_data = self._prepare_data(sorting_analyzer=sorting_analyzer, unit_ids=unit_ids, periods=periods) column_names_dtypes = {} for metric_name in metric_names: @@ -1605,7 +1618,7 @@ def _get_data(self, outputs="numpy", concatenated=False, return_data_name=None, self.sorting_analyzer.sorting, periods, ) - all_data = all_data[keep_mask] + all_data = slice_rows(all_data, keep_mask) # since we have the mask already, we can use it directly to avoid double computation spike_vector = self.sorting_analyzer.sorting.to_spike_vector(concatenated=True) sliced_spike_vector = spike_vector[keep_mask] @@ -1619,7 +1632,9 @@ def _get_data(self, outputs="numpy", concatenated=False, return_data_name=None, if outputs == "numpy": if copy and not self.sorting_analyzer._lazy: - return all_data.copy() # return a copy to avoid modification + # return a copy to avoid modification. `all_data` may be a zarr.Array even though this + # analyzer isn't itself lazy (e.g. shared by reference from a lazy source during merge/split) + return materialize_array(all_data) else: return all_data elif outputs == "by_unit": @@ -1637,7 +1652,7 @@ def _get_data(self, outputs="numpy", concatenated=False, return_data_name=None, data_by_units[segment_index] = {} for unit_id in unit_ids: inds = spike_indices[segment_index][unit_id] - data_by_units[segment_index][unit_id] = all_data[inds] + data_by_units[segment_index][unit_id] = slice_rows(all_data, inds) if concatenated: data_by_units_concatenated = { @@ -1659,7 +1674,7 @@ def _select_units_extension_data(self, unit_ids): new_data = dict() for data_name in self.nodepipeline_variables: if self.data.get(data_name) is not None: - new_data[data_name] = self.data[data_name][keep_spike_mask] + new_data[data_name] = slice_rows(self.data[data_name], keep_spike_mask) return new_data @@ -1670,9 +1685,10 @@ def _merge_extension_data( for data_name in self.nodepipeline_variables: if self.data.get(data_name) is not None: if keep_mask is None: - new_data[data_name] = self.data[data_name].copy() + # no filtering: sharing the reference is fine, materialization happens on save + new_data[data_name] = self.data[data_name] else: - new_data[data_name] = self.data[data_name][keep_mask] + new_data[data_name] = slice_rows(self.data[data_name], keep_mask) return new_data diff --git a/src/spikeinterface/core/core_tools.py b/src/spikeinterface/core/core_tools.py index ddc52e283f..17942b3ecd 100644 --- a/src/spikeinterface/core/core_tools.py +++ b/src/spikeinterface/core/core_tools.py @@ -782,24 +782,67 @@ def ms_to_samples(ms: float, sampling_frequency: float) -> int: return round(ms * sampling_frequency / 1000.0) -def slice_rows(array: np.ndarray | zarr.Array, row_indices: np.ndarray | list) -> np.ndarray: +def slice_rows(array: np.ndarray | zarr.Array, row_indices: np.ndarray | list, axis: int = 0) -> np.ndarray: """ - Slice a 2D array to select specific rows based on provided indices. + Slice an array to select specific indices/mask along one axis. Parameters ---------- array : np.ndarray | zarr.Array A numpy or zarr array or boolean mask from which rows will be selected. row_indices : np.ndarray | list - A list or array of row indices to select from the array. + A list or array of indices (or a boolean mask) to select along `axis`. + axis : int, default: 0 + The axis along which to select. Use 0 (the default) for spike/unit rows, + or a later axis (e.g. the channel axis of a 3D templates/waveforms array). Returns ------- np.ndarray - A new 2D numpy array containing only the selected rows. + A new numpy array containing only the selected entries along `axis`. """ if isinstance(array, zarr.Array): - # For zarr arrays, we need to convert the list of indices to a numpy array for advanced indexing - return array.oindex[row_indices] + # For zarr arrays, we need orthogonal indexing to select along a single axis + selection = (slice(None),) * axis + (row_indices,) + return array.oindex[selection] else: - return array[row_indices, ...] + selection = (slice(None),) * axis + (row_indices, ...) + return array[selection] + + +def materialize_array(array: np.ndarray | zarr.Array) -> np.ndarray: + """ + Return an independent, writable in-memory numpy array, ready for in-place mutation. + + Extension data is often shared by reference across analyzers (e.g. during merge/split) to + avoid unnecessary copies, since sharing is safe as long as nothing mutates the array in + place. A few operations (e.g. summing correlogram rows/columns into a merged unit, sparse + channel realignment of waveforms/PCA projections) do need to mutate in place, and for those + a real copy is required: zarr.Array has no `.copy()` method and is read-only from a + previously-saved store, while `copy.deepcopy()` does not actually clone a zarr array's + underlying data. Use this right before the in-place mutation, not as a default habit. + + Parameters + ---------- + array : np.ndarray | zarr.Array + A numpy or zarr array about to be mutated in place. + + Returns + ------- + np.ndarray + A new, independent, writable numpy array with the same data. + """ + if isinstance(array, zarr.Array): + return array[:] + else: + return array.copy() + + +def _ensure_seed(seed): + # when seed is None: + # we want to set one to push it in the Recordind._kwargs to reconstruct the same signal + # this is a better approach than having seed=42 or seed=my_dog_birthday because we ensure to have + # a new signal for all call with seed=None but the dump/load will still work + if seed is None: + seed = np.random.default_rng(seed=None).integers(0, 2**63) + return seed diff --git a/src/spikeinterface/core/generate.py b/src/spikeinterface/core/generate.py index 2c0ca2cfd5..28ae361fc0 100644 --- a/src/spikeinterface/core/generate.py +++ b/src/spikeinterface/core/generate.py @@ -12,17 +12,7 @@ from spikeinterface.core import BaseRecording, BaseRecordingSegment, BaseSorting from .snippets_tools import snippets_from_sorting -from .core_tools import define_function_from_class, ms_to_samples - - -def _ensure_seed(seed): - # when seed is None: - # we want to set one to push it in the Recordind._kwargs to reconstruct the same signal - # this is a better approach than having seed=42 or seed=my_dog_birthday because we ensure to have - # a new signal for all call with seed=None but the dump/load will still work - if seed is None: - seed = np.random.default_rng(seed=None).integers(0, 2**63) - return seed +from .core_tools import define_function_from_class, ms_to_samples, _ensure_seed def generate_recording( diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index c1b4d1f84f..3396b7dc88 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -53,6 +53,7 @@ def create_sorting_analyzer( format: Literal["memory", "binary_folder", "zarr"] = "memory", folder: str | Path | None = None, main_channel_indices: np.ndarray | None = None, + lazy: bool = False, peak_sign: PeakSignType = "both", peak_mode: PeakModeType = "extremum", num_spikes_for_main_channel: int = 100, @@ -264,6 +265,7 @@ def create_sorting_analyzer( recording=aggregated_recording, format=format, folder=folder, + lazy=lazy, sparse=sparse, sparsity=sparsity, main_channel_indices=main_channel_indices, @@ -346,6 +348,7 @@ def create_sorting_analyzer( recording, format=format, folder=folder, + lazy=lazy, main_channel_indices=main_channel_indices, peak_sign=peak_sign, peak_mode=peak_mode, @@ -358,7 +361,7 @@ def create_sorting_analyzer( def load_sorting_analyzer( - folder, load_extensions=True, format="auto", backend_options=None, lazy=False + folder, load_extensions=True, format="auto", backend_options=None, lazy=False, read_only=False ) -> "SortingAnalyzer": """ Load a SortingAnalyzer object from disk. @@ -380,6 +383,10 @@ def load_sorting_analyzer( * storage_options: dict | None (fsspec storage options) * saving_options: dict | None (additional saving options for creating and saving datasets) + lazy : bool, default: False + If True, the extensions are not loaded at load time, but only when they are accessed for the first time. + read_only : bool, default: False + If True, the SortingAnalyzer is loaded in read-only mode. This means that the extensions cannot be modified or deleted. Returns ------- @@ -388,7 +395,12 @@ def load_sorting_analyzer( """ return SortingAnalyzer.load( - folder, load_extensions=load_extensions, format=format, backend_options=backend_options, lazy=lazy + folder, + load_extensions=load_extensions, + format=format, + backend_options=backend_options, + lazy=lazy, + read_only=read_only, ) @@ -426,6 +438,7 @@ def __init__( peak_mode: PeakModeType = "extremum", backend_options: dict | None = None, lazy: bool = False, + read_only: bool = False, ): # very fast init because checks are done in load and create self.sorting = sorting @@ -456,6 +469,20 @@ def __init__( # the lazy flag is used to load the extensions in a lazy way (only when needed) self._lazy = lazy + # the read_only flag is used to load the extensions in a read-only way (cannot be modified) + self._read_only = read_only + + if self.format == "memory": + if self._lazy: + warnings.warn( + "Lazy mode is not supported for format='memory'. The extensions will be loaded in memory.", + ) + self._lazy = False + if self._read_only: + warnings.warn( + "Read-only mode is not supported for format='memory'. Extensions can be modified in memory, but changes will not be saved to disk.", + ) + self._read_only = False # extensions are not loaded at init self.extensions = dict() @@ -497,6 +524,7 @@ def create( main_channel_indices: np.ndarray, format: Literal["memory", "binary_folder", "zarr"] = "memory", folder: str | Path | None = None, + lazy: bool = False, sparsity: ChannelSparsity | None = None, return_scaled: bool | None = None, return_in_uV: bool = True, @@ -565,6 +593,7 @@ def create( peak_mode, rec_attributes=None, backend_options=backend_options, + lazy=lazy, ) elif format == "zarr": assert folder is not None, "For format='zarr' folder must be provided" @@ -578,6 +607,7 @@ def create( peak_mode, rec_attributes=None, backend_options=backend_options, + lazy=lazy, ) else: raise ValueError(f"SortingAnalyzer.create: wrong format {format}") @@ -597,6 +627,7 @@ def load( format: Literal["auto", "binary_folder", "zarr"] = "auto", backend_options: dict | None = None, lazy: bool = False, + read_only: bool = False, ): """ Load folder or zarr. @@ -610,11 +641,11 @@ def load( if format == "binary_folder": sorting_analyzer = SortingAnalyzer.load_from_binary_folder( - folder, recording=recording, backend_options=backend_options, lazy=lazy + folder, recording=recording, backend_options=backend_options, lazy=lazy, read_only=read_only ) elif format == "zarr": sorting_analyzer = SortingAnalyzer.load_from_zarr( - folder, recording=recording, backend_options=backend_options, lazy=lazy + folder, recording=recording, backend_options=backend_options, lazy=lazy, read_only=read_only ) else: raise ValueError(f"SortingAnalyzer.load: wrong format {format}") @@ -679,6 +710,7 @@ def create_binary_folder( peak_mode: PeakModeType, rec_attributes: dict | None, backend_options: dict | None, + lazy: bool = False, ) -> "SortingAnalyzer": # used by create and save_as folder = Path(folder) @@ -756,7 +788,7 @@ def create_binary_folder( if probegroup is not None: probeinterface.write_probeinterface(probegroup_file, probegroup) - return cls.load_from_binary_folder(folder, recording=recording, backend_options=backend_options) + return cls.load_from_binary_folder(folder, recording=recording, backend_options=backend_options, lazy=lazy) @classmethod def _handle_backward_compatibility_settings_pre_init(cls, settings: dict[str, Any]): @@ -902,6 +934,7 @@ def load_from_binary_folder( recording: BaseRecording | None = None, backend_options: dict | None = None, lazy: bool = False, + read_only: bool = False, ) -> "SortingAnalyzer": from .loading import load @@ -997,6 +1030,7 @@ def load_from_binary_folder( peak_mode=settings["peak_mode"], backend_options=backend_options, lazy=lazy, + read_only=read_only, ) sorting_analyzer.folder = folder @@ -1021,6 +1055,7 @@ def create_zarr( peak_mode: PeakModeType, rec_attributes: dict | None, backend_options: dict | None, + lazy: bool = False, ) -> "SortingAnalyzer": # used by create and save_as import zarr @@ -1107,7 +1142,7 @@ def create_zarr( # Consolidate metadata (for faster reads) zarr.consolidate_metadata(zarr_root.store) - return cls.load_from_zarr(folder, recording=recording, backend_options=backend_options) + return cls.load_from_zarr(folder, recording=recording, backend_options=backend_options, lazy=lazy) @classmethod def load_from_zarr( @@ -1116,6 +1151,7 @@ def load_from_zarr( recording: BaseRecording | None = None, backend_options: dict | None = None, lazy: bool = False, + read_only: bool = False, ) -> "SortingAnalyzer": import zarr from .loading import load @@ -1201,6 +1237,7 @@ def load_from_zarr( peak_mode=settings["peak_mode"], backend_options=backend_options, lazy=lazy, + read_only=read_only, ) sorting_analyzer.folder = folder @@ -1526,11 +1563,17 @@ def _save_or_select_or_merge_or_split( new_sorting_analyzer : SortingAnalyzer The newly created SortingAnalyzer object. """ - if self._lazy: + if self._read_only: raise ValueError( - "Cannot save, select, merge or split units when the SortingAnalyzer is lazy. " - "Please load the SortingAnalyzer with lazy=False." + "Cannot save, select, merge or split units when the SortingAnalyzer is read-only. " + "Please load the SortingAnalyzer with read_only=False." ) + if self._lazy: + # extensions are only registered in self.extensions on first access when lazy, so a lazily + # loaded analyzer that hasn't touched every extension yet would otherwise silently lose any + # untouched extension during copy/select/merge/split below. Force-load the registry (this + # keeps each extension's data as a memmap/zarr handle, it does not eagerly materialize it). + self.load_all_saved_extension() if self.has_recording(): recording = self._recording elif self.has_temporary_recording(): @@ -1658,6 +1701,7 @@ def _save_or_select_or_merge_or_split( self.peak_mode, self.rec_attributes, backend_options=backend_options, + lazy=self._lazy, ) elif format == "zarr": @@ -1673,6 +1717,7 @@ def _save_or_select_or_merge_or_split( self.peak_mode, self.rec_attributes, backend_options=backend_options, + lazy=self._lazy, ) else: raise ValueError(f"SortingAnalyzer.save: unsupported format: {format}") @@ -2061,6 +2106,8 @@ def copy(self): return self._save_or_select_or_merge_or_split(format="memory", folder=None) def is_read_only(self) -> bool: + if self._read_only: + return True if self.format == "memory": return False elif self.format == "binary_folder": @@ -2224,7 +2271,7 @@ def compute(self, input, save=True, extension_params=None, verbose=False, **kwar * a list: compute several extensions. The list contains the extension names. Additional parameters can be passed with the extension_params argument. save : bool, default: True - If True the extension is saved to disk (only if sorting analyzer format is not "memory") + If True the extension is saved to disk (only if sorting analyzer format is not "memory"). extension_params : dict or None, default: None If input is a list, this parameter can be used to specify parameters for each extension. The extension_params keys must be included in the input list. @@ -2257,9 +2304,9 @@ 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. + if self._read_only: + # If the analyzer is read-only, 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 read-only. save = False if isinstance(input, str): return self.compute_one_extension(extension_name=input, save=save, verbose=verbose, **kwargs) @@ -2300,8 +2347,8 @@ def compute_one_extension(self, extension_name, save=True, verbose=False, **kwar The name of the extension. For instance "waveforms", "templates", ... save : bool, default: True - It the extension can be saved then it is saved. - If not then the extension will only live in memory as long as the object is deleted. + If True the extension is saved to disk (only if sorting analyzer format is not "memory"). + If False the extension will only live in memory as long as the object is deleted. save=False is convenient to try some parameters without changing an already saved extension. **kwargs: @@ -2323,6 +2370,8 @@ def compute_one_extension(self, extension_name, save=True, verbose=False, **kwar >>> wfs = compute_waveforms(sorting_analyzer, **some_params) """ + if self._read_only: + save = False extension_class = get_extension_class(extension_name) for child in _get_children_dependencies(extension_name): @@ -2372,6 +2421,7 @@ def compute_several_extensions(self, extensions, save=True, verbose=False, **job It the extension can be saved then it is saved. If not then the extension will only live in memory as long as the object is deleted. save=False is convenient to try some parameters without changing an already saved extension. + If True the extension is saved to disk (only if sorting analyzer format is not "memory"). Returns ------- @@ -2384,6 +2434,8 @@ def compute_several_extensions(self, extensions, save=True, verbose=False, **job >>> sorting_analyzer.compute_several_extensions({"waveforms": {"ms_before": 1.2}, "templates" : {"operators": ["average", "std"]}}) """ + if self._read_only: + save = False # Check dependencies: either already computed or in the extensions to compute extensions_to_compute = list(extensions.keys()) for extension_name, extension_params in extensions.items(): @@ -2574,9 +2626,13 @@ def delete_extension(self, extension_name) -> None: """ Delete the extension from the dict and also in the persistent zarr or folder. """ - # delete from folder or zarr - if self.format != "memory" and self.has_extension(extension_name) and not self._lazy: + if self.format != "memory" and self.has_extension(extension_name) and not self._read_only: + # close any memmap handles already held (e.g. by a previous lazy load), possibly shared with + # an external reference, before touching disk (ext.delete() below closes its own) + old_extension = self.extensions.get(extension_name) + if old_extension is not None: + old_extension._close_memmaps() # need a reload to reset the folder ext = self.load_extension(extension_name) ext.delete() @@ -3433,6 +3489,9 @@ def delete(self): """ Delete the extension from the folder or zarr and from the dict. """ + # close any memmap handles onto these files first (e.g. from a lazy load) so the + # folder/files can actually be removed + self._close_memmaps() self._delete_extension_folder() self.params = None self.run_info = self._default_run_info_dict() @@ -3448,6 +3507,24 @@ def reset(self): self.run_info = self._default_run_info_dict() self.data = dict() + def _close_memmaps(self): + """ + Close any open np.memmap handles held by this extension's data. + + This must run before this extension's on-disk files are deleted/overwritten (e.g. when + recomputing an extension that a previous, possibly lazily-loaded, instance is still + registered for): on Windows a memory-mapped file cannot be deleted or rewritten while a + handle to it is still open. Only memmap-backed entries are touched: other already + materialized data (e.g. a DataFrame of previously computed metrics, used to carry forward + results that are not being recomputed) is left untouched. + """ + for key, value in list(self.data.items()): + if isinstance(value, np.memmap): + mmap_obj = getattr(value, "_mmap", None) + if mmap_obj is not None: + mmap_obj.close() + self.data[key] = None + def set_params(self, save=True, **params): """ Set parameters for the extension and @@ -3456,6 +3533,11 @@ def set_params(self, save=True, **params): # this ensure data is also deleted and corresponds to params # this also ensure the group is created if save: + # if a previous instance of this extension is still registered (e.g. loaded lazily), + # close any memmap handles it holds before we delete/overwrite its on-disk files + old_extension = self.sorting_analyzer.extensions.get(self.extension_name) + if old_extension is not None and old_extension is not self: + old_extension._close_memmaps() self._reset_extension_folder() params = self._set_params(**params) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 25aeb78c1a..dc6a5a1311 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -400,12 +400,18 @@ def test_load_in_lazy_mode(tmp_path, dataset, format): if isinstance(value, np.ndarray): assert isinstance(value, array_class) - # check that the lazy mode does not overwrite existing extensions + # a lazy (but not read-only) analyzer is allowed to 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 + assert random_spikes_ext.params["max_spikes_per_unit"] == 10 + + # check that a lazy+read-only analyzer does not overwrite existing extensions + sorting_analyzer_lazy_ro = load_sorting_analyzer(folder, format="auto", lazy=True, read_only=True) + sorting_analyzer_lazy_ro.compute("random_spikes", max_spikes_per_unit=20) + 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"] != 20 def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder): diff --git a/src/spikeinterface/generation/drifting_generator.py b/src/spikeinterface/generation/drifting_generator.py index 21d25cdd41..f217dda1b3 100644 --- a/src/spikeinterface/generation/drifting_generator.py +++ b/src/spikeinterface/generation/drifting_generator.py @@ -14,13 +14,13 @@ from spikeinterface import Templates from spikeinterface.core import ms_to_samples +from spikeinterface.core.core_tools import _ensure_seed from spikeinterface.core.generate import ( generate_unit_locations, generate_sorting, generate_templates, synthesize_amplitude_factor, _ensure_unit_params, - _ensure_seed, ) from .drift_tools import DriftingTemplates, make_linear_displacement, InjectDriftingTemplatesRecording from .noise_tools import generate_noise diff --git a/src/spikeinterface/generation/hybrid_tools.py b/src/spikeinterface/generation/hybrid_tools.py index 84e6187ac7..426699f743 100644 --- a/src/spikeinterface/generation/hybrid_tools.py +++ b/src/spikeinterface/generation/hybrid_tools.py @@ -4,12 +4,12 @@ import numpy as np from spikeinterface.core import BaseRecording, BaseSorting, Templates, ms_to_samples +from spikeinterface.core.core_tools import _ensure_seed from spikeinterface.core.generate import ( generate_templates, generate_unit_locations, generate_sorting, InjectTemplatesRecording, - _ensure_seed, synthesize_amplitude_factor, ) diff --git a/src/spikeinterface/generation/noise_tools.py b/src/spikeinterface/generation/noise_tools.py index 87c9c52387..7fbecf0c9b 100644 --- a/src/spikeinterface/generation/noise_tools.py +++ b/src/spikeinterface/generation/noise_tools.py @@ -3,8 +3,7 @@ import numpy as np from spikeinterface.core import BaseRecording, BaseRecordingSegment -from spikeinterface.core.generate import _ensure_seed -from spikeinterface.core.core_tools import define_function_from_class +from spikeinterface.core.core_tools import define_function_from_class, _ensure_seed class NoiseGeneratorRecording(BaseRecording): diff --git a/src/spikeinterface/metrics/conftest.py b/src/spikeinterface/metrics/conftest.py index 11d06c636a..899586dd84 100644 --- a/src/spikeinterface/metrics/conftest.py +++ b/src/spikeinterface/metrics/conftest.py @@ -21,7 +21,7 @@ def make_small_analyzer(): extensions_to_compute = { "random_spikes": {"seed": 1205}, - "noise_levels": {"seed": 1205}, + "noise_levels": {"random_slices_kwargs": {"seed": 1205}}, "waveforms": {}, "templates": {"operators": ["average", "median"]}, "spike_amplitudes": {}, diff --git a/src/spikeinterface/metrics/quality/misc_metrics.py b/src/spikeinterface/metrics/quality/misc_metrics.py index c8b8f33a08..3cd0be6b1f 100644 --- a/src/spikeinterface/metrics/quality/misc_metrics.py +++ b/src/spikeinterface/metrics/quality/misc_metrics.py @@ -34,6 +34,69 @@ HAVE_NUMBA = False +# Metrics that read spike_amplitudes/amplitude_scalings by unit. Their by-unit amplitude data is +# pre-fetched once in `ComputeQualityMetrics._prepare_data` and shared via `tmp_data` instead of each +# metric independently re-fetching the full array from the extension. +amplitude_based_metric_names = {"amplitude_cutoff", "amplitude_median", "noise_cutoff", "amplitude_cv", "sd_ratio"} + + +def _amplitude_extension_name_for_metric(metric_name, metric_params, sorting_analyzer): + if metric_name == "amplitude_median": + # only depends on spike_amplitudes (no amplitude_scalings fallback) + return "spike_amplitudes" + if metric_name == "sd_ratio": + return "spike_amplitudes" + if metric_name == "amplitude_cv": + configured = metric_params.get("amplitude_cv", {}).get("amplitude_extension") + if configured is not None: + return configured + return "spike_amplitudes" if sorting_analyzer.has_extension("spike_amplitudes") else "amplitude_scalings" + if metric_name in ("amplitude_cutoff", "noise_cutoff"): + return "spike_amplitudes" if sorting_analyzer.has_extension("spike_amplitudes") else "amplitude_scalings" + raise ValueError(f"Unknown amplitude-based metric name: {metric_name}") + + +def get_amplitudes_by_segment(sorting_analyzer, requested_metric_names, metric_params, periods=None): + """Fetch each needed spike_amplitudes/amplitude_scalings extension's by-unit data (per-segment, + unconcatenated) once, so it can be shared across all requested amplitude-based metrics.""" + needed_extension_names = { + _amplitude_extension_name_for_metric(metric_name, metric_params, sorting_analyzer) + for metric_name in requested_metric_names + } + + amplitudes_by_segment = {} + for extension_name in needed_extension_names: + extension = sorting_analyzer.get_extension(extension_name) + if extension is not None: + amplitudes_by_segment[extension_name] = extension.get_data( + outputs="by_unit", concatenated=False, periods=periods + ) + return amplitudes_by_segment + + +def _amplitudes_by_unit_from_cache(sorting_analyzer, extension_name, periods, concatenated, tmp_data): + """Return by-unit amplitude data for `extension_name`, reusing the pre-fetched `tmp_data` cache + (built by `get_amplitudes_by_segment`) when available, falling back to a direct fetch otherwise + (e.g. when the metric function is called directly, outside the metrics-dataframe machinery).""" + amplitudes_by_segment = None + if tmp_data is not None: + amplitudes_by_segment = tmp_data.get("amplitudes_by_segment", {}).get(extension_name) + if amplitudes_by_segment is None: + extension = sorting_analyzer.get_extension(extension_name) + amplitudes_by_segment = extension.get_data(outputs="by_unit", concatenated=False, periods=periods) + + if not concatenated: + return amplitudes_by_segment + + num_segments = len(amplitudes_by_segment) + return { + unit_id: np.concatenate( + [amplitudes_by_segment[segment_index][unit_id] for segment_index in range(num_segments)] + ) + for unit_id in sorting_analyzer.unit_ids + } + + def compute_presence_ratios( sorting_analyzer, unit_ids=None, periods=None, bin_duration_s=60.0, mean_fr_ratio_thresh=0.0 ): @@ -805,11 +868,12 @@ class FiringRange(BaseMetric): def compute_amplitude_cv_metrics( sorting_analyzer, unit_ids=None, + tmp_data=None, periods=None, average_num_spikes_per_bin=50, percentiles=(5, 95), min_num_bins=10, - amplitude_extension="spike_amplitudes", + amplitude_extension=None, ): """ Calculate coefficient of variation of spike amplitudes within defined temporal bins. @@ -834,8 +898,9 @@ def compute_amplitude_cv_metrics( min_num_bins : int, default: 10 The minimum number of bins to compute the median and range. If the number of bins is less than this then the median and range are set to NaN. - amplitude_extension : str, default: "spike_amplitudes" - The name of the extension to load the amplitudes from. "spike_amplitudes" or "amplitude_scalings". + amplitude_extension : "spike_amplitudes" | "amplitude_scalings" | None, default: None + The name of the extension to load the amplitudes from. If None, "spike_amplitudes" is used + if available, otherwise "amplitude_scalings". Returns ------- @@ -850,6 +915,10 @@ def compute_amplitude_cv_metrics( """ check_has_required_extensions("amplitude_cv", sorting_analyzer) res = namedtuple("amplitude_cv", ["amplitude_cv_median", "amplitude_cv_range"]) + if amplitude_extension is None: + amplitude_extension = ( + "spike_amplitudes" if sorting_analyzer.has_extension("spike_amplitudes") else "amplitude_scalings" + ) assert amplitude_extension in ( "spike_amplitudes", "amplitude_scalings", @@ -861,8 +930,8 @@ def compute_amplitude_cv_metrics( total_durations = compute_total_durations_per_unit(sorting_analyzer, periods=periods) num_spikes = sorting.count_num_spikes_per_unit(outputs="dict", unit_ids=unit_ids) - amps = sorting_analyzer.get_extension(amplitude_extension).get_data( - outputs="by_unit", concatenated=False, periods=periods + amps = _amplitudes_by_unit_from_cache( + sorting_analyzer, amplitude_extension, periods, concatenated=False, tmp_data=tmp_data ) amplitude_cv_medians, amplitude_cv_ranges = {}, {} @@ -909,7 +978,7 @@ class AmplitudeCV(BaseMetric): "average_num_spikes_per_bin": 50, "percentiles": (5, 95), "min_num_bins": 10, - "amplitude_extension": "spike_amplitudes", + "amplitude_extension": None, } metric_columns = {"amplitude_cv_median": float, "amplitude_cv_range": float} metric_descriptions = { @@ -917,12 +986,14 @@ class AmplitudeCV(BaseMetric): "amplitude_cv_range": "Range of the coefficient of variation of spike amplitudes within temporal bins.", } supports_periods = True + needs_tmp_data = True depend_on = ["spike_amplitudes|amplitude_scalings"] def compute_amplitude_cutoffs( sorting_analyzer, unit_ids=None, + tmp_data=None, periods=None, num_histogram_bins=500, histogram_smoothing_value=3, @@ -977,8 +1048,9 @@ def compute_amplitude_cutoffs( available_extension = ( "spike_amplitudes" if sorting_analyzer.has_extension("spike_amplitudes") else "amplitude_scalings" ) - extension = sorting_analyzer.get_extension(available_extension) - amplitudes_by_units = extension.get_data(outputs="by_unit", concatenated=True, periods=periods) + amplitudes_by_units = _amplitudes_by_unit_from_cache( + sorting_analyzer, available_extension, periods, concatenated=True, tmp_data=tmp_data + ) for unit_id in unit_ids: amplitudes = amplitudes_by_units[unit_id] @@ -1015,10 +1087,11 @@ class AmplitudeCutoff(BaseMetric): "amplitude_cutoff": "Estimated fraction of missing spikes, based on the amplitude distribution." } supports_periods = True + needs_tmp_data = True depend_on = ["spike_amplitudes|amplitude_scalings"] -def compute_amplitude_medians(sorting_analyzer, unit_ids=None, periods=None): +def compute_amplitude_medians(sorting_analyzer, unit_ids=None, tmp_data=None, periods=None): """ Compute median of the amplitude distributions. @@ -1048,8 +1121,9 @@ def compute_amplitude_medians(sorting_analyzer, unit_ids=None, periods=None): unit_ids = sorting_analyzer.unit_ids all_amplitude_medians = {} - amplitude_extension = sorting_analyzer.get_extension("spike_amplitudes") - amplitudes_by_units = amplitude_extension.get_data(outputs="by_unit", concatenated=True, periods=periods) + amplitudes_by_units = _amplitudes_by_unit_from_cache( + sorting_analyzer, "spike_amplitudes", periods, concatenated=True, tmp_data=tmp_data + ) for unit_id in unit_ids: all_amplitude_medians[unit_id] = np.median(amplitudes_by_units[unit_id]) @@ -1062,11 +1136,12 @@ class AmplitudeMedian(BaseMetric): metric_columns = {"amplitude_median": float} metric_descriptions = {"amplitude_median": "Median of the amplitude distributions for each unit in µV."} supports_periods = True + needs_tmp_data = True depend_on = ["spike_amplitudes"] def compute_noise_cutoffs( - sorting_analyzer, unit_ids=None, periods=None, high_quantile=0.25, low_quantile=0.1, n_bins=100 + sorting_analyzer, unit_ids=None, tmp_data=None, periods=None, high_quantile=0.25, low_quantile=0.1, n_bins=100 ): """ A metric to determine if a unit's amplitude distribution is cut off as it approaches zero, without assuming a Gaussian distribution. @@ -1116,8 +1191,9 @@ def compute_noise_cutoffs( available_extension = ( "spike_amplitudes" if sorting_analyzer.has_extension("spike_amplitudes") else "amplitude_scalings" ) - extension = sorting_analyzer.get_extension(available_extension) - amplitudes_by_units = extension.get_data(outputs="by_unit", concatenated=True, periods=periods) + amplitudes_by_units = _amplitudes_by_unit_from_cache( + sorting_analyzer, available_extension, periods, concatenated=True, tmp_data=tmp_data + ) for unit_id in unit_ids: amplitudes = amplitudes_by_units[unit_id] @@ -1148,6 +1224,7 @@ class NoiseCutoff(BaseMetric): "noise_ratio": "Ratio of counts in the lower-amplitude bins to the count in the highest bin.", } supports_periods = True + needs_tmp_data = True depend_on = ["spike_amplitudes|amplitude_scalings"] @@ -1228,10 +1305,9 @@ def compute_drift_metrics( spike_locations_by_unit_and_segments = spike_locations_ext.get_data( outputs="by_unit", concatenated=False, periods=periods ) - spike_locations_by_unit = spike_locations_ext.get_data(outputs="by_unit", concatenated=True, periods=periods) segment_samples = [sorting_analyzer.get_num_samples(i) for i in range(sorting_analyzer.get_num_segments())] - data = spike_locations_by_unit[unit_ids[0]] + data = spike_locations_by_unit_and_segments[0][unit_ids[0]] assert direction in data.dtype.names, ( f"Direction {direction} is invalid. Available directions: " f"{data.dtype.names}" ) @@ -1249,8 +1325,15 @@ def compute_drift_metrics( reference_positions = {} median_position_segments = {unit_id: np.array([]) for unit_id in unit_ids} + num_segments = sorting_analyzer.get_num_segments() for unit_id in unit_ids: - reference_positions[unit_id] = np.median(spike_locations_by_unit[unit_id][direction]) + unit_direction_values = np.concatenate( + [ + spike_locations_by_unit_and_segments[segment_index][unit_id][direction] + for segment_index in range(num_segments) + ] + ) + reference_positions[unit_id] = np.median(unit_direction_values) for segment_index in range(sorting_analyzer.get_num_segments()): for unit_id in unit_ids: @@ -1326,6 +1409,7 @@ class Drift(BaseMetric): def compute_sd_ratio( sorting_analyzer: SortingAnalyzer, unit_ids=None, + tmp_data=None, periods=None, censored_period_ms: float = 4.0, correct_for_drift: bool = True, @@ -1385,8 +1469,8 @@ def compute_sd_ratio( ) return {unit_id: np.nan for unit_id in unit_ids} - spike_amplitudes = sorting_analyzer.get_extension("spike_amplitudes").get_data( - outputs="by_unit", concatenated=False, periods=periods + spike_amplitudes = _amplitudes_by_unit_from_cache( + sorting_analyzer, "spike_amplitudes", periods, concatenated=False, tmp_data=tmp_data ) if not HAVE_NUMBA: @@ -1472,6 +1556,7 @@ class SDRatio(BaseMetric): } needs_recording = True supports_periods = True + needs_tmp_data = True depend_on = ["templates", "spike_amplitudes"] diff --git a/src/spikeinterface/metrics/quality/pca_metrics.py b/src/spikeinterface/metrics/quality/pca_metrics.py index 44d7eb8954..45df10a833 100644 --- a/src/spikeinterface/metrics/quality/pca_metrics.py +++ b/src/spikeinterface/metrics/quality/pca_metrics.py @@ -140,9 +140,9 @@ class NearestNeighbor(BaseMetric): def _nn_advanced_one_unit(args): - unit_id, sorting_analyzer_or_folder, n_spikes_all_units, fr_all_units, metric_params, seed = args + unit_id, sorting_analyzer_or_folder, n_spikes_all_units, fr_all_units, metric_params, seed, lazy = args if isinstance(sorting_analyzer_or_folder, (str, Path)): - sorting_analyzer = load(sorting_analyzer_or_folder) + sorting_analyzer = load(sorting_analyzer_or_folder, lazy=lazy) else: sorting_analyzer = sorting_analyzer_or_folder @@ -229,7 +229,15 @@ def _nn_advanced_metric_function(sorting_analyzer, unit_ids, tmp_data, job_kwarg for unit_id in units_loop: _, nn_isolation, nn_unit_id, nn_noise_overlap = _nn_advanced_one_unit( - (unit_id, sorting_analyzer, n_spikes_all_units, fr_all_units, metric_params, seed) + ( + unit_id, + sorting_analyzer, + n_spikes_all_units, + fr_all_units, + metric_params, + seed, + sorting_analyzer._lazy, + ) ) nn_isolation_dict[unit_id] = nn_isolation nn_noise_overlap_dict[unit_id] = nn_noise_overlap @@ -243,7 +251,17 @@ def _nn_advanced_metric_function(sorting_analyzer, unit_ids, tmp_data, job_kwarg # If we got here, we are sure the sorting_analyzer is saved on disk args_list = [] for unit_id in unit_ids: - args_list.append((unit_id, sorting_analyzer.folder, n_spikes_all_units, fr_all_units, metric_params, seed)) + args_list.append( + ( + unit_id, + sorting_analyzer.folder, + n_spikes_all_units, + fr_all_units, + metric_params, + seed, + sorting_analyzer._lazy, + ) + ) with ProcessPoolExecutor( max_workers=n_jobs, diff --git a/src/spikeinterface/metrics/quality/quality_metrics.py b/src/spikeinterface/metrics/quality/quality_metrics.py index d40ed63cdf..8f5f98b376 100644 --- a/src/spikeinterface/metrics/quality/quality_metrics.py +++ b/src/spikeinterface/metrics/quality/quality_metrics.py @@ -137,13 +137,29 @@ def _set_params( skip_pc_metrics=skip_pc_metrics, ) - def _prepare_data(self, sorting_analyzer, unit_ids=None): + def _prepare_data(self, sorting_analyzer, unit_ids=None, periods=None): """Prepare shared data for quality metrics computation.""" # Pre-compute shared PCA data from spikeinterface.metrics.spiketrain.metrics import compute_num_spikes, compute_firing_rates + from .misc_metrics import amplitude_based_metric_names, get_amplitudes_by_segment tmp_data = {} + if unit_ids is None: + unit_ids = sorting_analyzer.unit_ids + + # Pre-fetch amplitude data (spike_amplitudes/amplitude_scalings) once and share it across + # amplitude_cutoff/amplitude_median/noise_cutoff/amplitude_cv/sd_ratio, instead of each metric + # independently rebuilding its own full by-unit copy of the same underlying array. + requested_amplitude_metrics = [m for m in self.params["metric_names"] if m in amplitude_based_metric_names] + if requested_amplitude_metrics: + tmp_data["amplitudes_by_segment"] = get_amplitudes_by_segment( + sorting_analyzer, + requested_amplitude_metrics, + self.params["metric_params"], + periods=periods, + ) + # Check if any PCA metrics are requested pca_metric_names = [m.metric_name for m in pca_metrics_list] requested_pca_metrics = [m for m in self.params["metric_names"] if m in pca_metric_names] @@ -156,9 +172,6 @@ def _prepare_data(self, sorting_analyzer, unit_ids=None): if pca_ext is None: return tmp_data - if unit_ids is None: - unit_ids = sorting_analyzer.unit_ids - # Get dense PCA projections for all requested units dense_projections, spike_unit_indices = pca_ext.get_some_projections(channel_ids=None, unit_ids=unit_ids) all_labels = sorting_analyzer.sorting.unit_ids[spike_unit_indices] diff --git a/src/spikeinterface/metrics/quality/tests/test_metrics_functions.py b/src/spikeinterface/metrics/quality/tests/test_metrics_functions.py index b52673d9fd..b69411164a 100644 --- a/src/spikeinterface/metrics/quality/tests/test_metrics_functions.py +++ b/src/spikeinterface/metrics/quality/tests/test_metrics_functions.py @@ -790,7 +790,7 @@ def test_unit_id_order_independence(small_sorting_analyzer): extensions_to_compute = { "random_spikes": {"seed": 1205}, - "noise_levels": {"seed": 1205}, + "noise_levels": {"random_slices_kwargs": {"seed": 1205}}, "waveforms": {}, "templates": {}, "spike_amplitudes": {}, diff --git a/src/spikeinterface/metrics/quality/tests/test_pca_metrics.py b/src/spikeinterface/metrics/quality/tests/test_pca_metrics.py index 85d1c331d5..3f7dc396ea 100644 --- a/src/spikeinterface/metrics/quality/tests/test_pca_metrics.py +++ b/src/spikeinterface/metrics/quality/tests/test_pca_metrics.py @@ -57,6 +57,67 @@ def test_compute_pc_metrics_multi_processing(small_sorting_analyzer, tmp_path): assert np.array_equal(values1, values2) +def test_compute_pc_metrics_lazy_loading(small_sorting_analyzer, tmp_path): + """ + PCA metrics (including the multi-processed nn_advanced metric) must give the same + results whether the SortingAnalyzer is loaded eagerly (extension data fully read into + RAM) or lazily (extension data kept as memmap/zarr handles). This is the cheapest way + to reduce baseline memory usage for quality metrics on very large recordings. + """ + sorting_analyzer = small_sorting_analyzer + metric_names = get_quality_pca_metric_list() + metric_params = dict(nn_advanced=dict(seed=2308)) + + sorting_analyzer_saved = sorting_analyzer.save_as(folder=tmp_path / "analyzer_lazy", format="binary_folder") + + res_eager = compute_quality_metrics( + sorting_analyzer_saved, + metric_names=metric_names, + n_jobs=2, + seed=1205, + metric_params=metric_params, + ) + + from spikeinterface.core import load_sorting_analyzer + + sorting_analyzer_lazy = load_sorting_analyzer(tmp_path / "analyzer_lazy", format="auto", lazy=True) + res_lazy = compute_quality_metrics( + sorting_analyzer_lazy, + metric_names=metric_names, + n_jobs=2, + seed=1205, + metric_params=metric_params, + ) + + for metric_name in res_eager.columns: + values_eager = res_eager[metric_name].values + values_lazy = res_lazy[metric_name].values + + if values_eager.dtype.kind == "f": + np.testing.assert_almost_equal(values_eager, values_lazy, decimal=4) + else: + assert np.array_equal(values_eager, values_lazy) + + # a lazy (but not read-only) analyzer is allowed to persist to disk by default. + # the eager computation above already did this (analyzer is not lazy), so reset first. + sorting_analyzer_saved.delete_extension("quality_metrics") + sorting_analyzer_lazy2 = load_sorting_analyzer(tmp_path / "analyzer_lazy", format="auto", lazy=True) + sorting_analyzer_lazy2.compute("quality_metrics", metric_names=metric_names, seed=1205, metric_params=metric_params) + reloaded = load_sorting_analyzer(tmp_path / "analyzer_lazy", format="auto", lazy=True) + assert reloaded.has_extension("quality_metrics") + + # a lazy AND read-only analyzer must not persist to disk, even if asked to explicitly + sorting_analyzer_saved.delete_extension("quality_metrics") + sorting_analyzer_lazy_ro = load_sorting_analyzer( + tmp_path / "analyzer_lazy", format="auto", lazy=True, read_only=True + ) + sorting_analyzer_lazy_ro.compute( + "quality_metrics", metric_names=metric_names, seed=1205, metric_params=metric_params, save=True + ) + reloaded = load_sorting_analyzer(tmp_path / "analyzer_lazy", format="auto", lazy=True) + assert not reloaded.has_extension("quality_metrics") + + if __name__ == "__main__": from spikeinterface.metrics.conftest import make_small_analyzer diff --git a/src/spikeinterface/metrics/template/template_metrics.py b/src/spikeinterface/metrics/template/template_metrics.py index 97c9c8cbad..4e843e6726 100644 --- a/src/spikeinterface/metrics/template/template_metrics.py +++ b/src/spikeinterface/metrics/template/template_metrics.py @@ -228,7 +228,7 @@ def _set_params( min_extremum_distance_samples=min_extremum_distance_samples, ) - def _prepare_data(self, sorting_analyzer, unit_ids): + def _prepare_data(self, sorting_analyzer, unit_ids, periods=None): import warnings import pandas as pd from scipy.signal import resample_poly diff --git a/src/spikeinterface/postprocessing/correlograms.py b/src/spikeinterface/postprocessing/correlograms.py index 13431a8dbb..0c9c5a7aa9 100644 --- a/src/spikeinterface/postprocessing/correlograms.py +++ b/src/spikeinterface/postprocessing/correlograms.py @@ -1,7 +1,6 @@ import importlib.util import warnings import platform -from copy import deepcopy from tqdm.auto import tqdm from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor @@ -13,6 +12,7 @@ from spikeinterface.core import BaseSorting from spikeinterface.core.job_tools import fix_job_kwargs, _shared_job_kwargs_doc +from spikeinterface.core.core_tools import slice_rows, materialize_array from spikeinterface.core.sortinganalyzer import ( AnalyzerExtension, SortingAnalyzer, @@ -99,7 +99,7 @@ def _set_params(self, window_ms: float = 50.0, bin_ms: float = 1.0, method: str def _select_units_extension_data(self, unit_ids): # filter metrics dataframe unit_indices = self.sorting_analyzer.sorting.ids_to_indices(unit_ids) - new_ccgs = self.data["ccgs"][unit_indices][:, unit_indices] + new_ccgs = slice_rows(self.data["ccgs"], unit_indices)[:, unit_indices] new_bins = self.data["bins"] new_data = dict(ccgs=new_ccgs, bins=new_bins) return new_data @@ -166,7 +166,8 @@ def _merge_extension_data( if unit_involved_in_merge is False: old_to_new_unit_index_map[old_unit_index] = new_sorting_analyzer.sorting.id_to_index(old_unit) - correlograms, new_bins = deepcopy(self.get_data()) + correlograms = materialize_array(self.data["ccgs"]) + new_bins = self.data["bins"] for new_unit_id, merge_unit_group in zip(new_unit_ids, merge_unit_groups): merge_unit_group_indices = self.sorting_analyzer.sorting.ids_to_indices(merge_unit_group) @@ -274,7 +275,7 @@ def _set_params(self, window_ms: float = 50.0, bin_ms: float = 1.0, method: str def _select_units_extension_data(self, unit_ids): # filter metrics dataframe unit_indices = self.sorting_analyzer.sorting.ids_to_indices(unit_ids) - new_acgs = self.data["acgs"][unit_indices] + new_acgs = slice_rows(self.data["acgs"], unit_indices) new_bins = self.data["bins"] new_data = dict(ccgs=new_acgs, bins=new_bins) return new_data @@ -1212,8 +1213,8 @@ def _set_params( def _select_units_extension_data(self, unit_ids): # filter metrics dataframe unit_indices = self.sorting_analyzer.sorting.ids_to_indices(unit_ids) - new_acgs_3d = self.data["acgs_3d"][unit_indices] - new_firing_quantiles = self.data["firing_quantiles"][unit_indices] + new_acgs_3d = slice_rows(self.data["acgs_3d"], unit_indices) + new_firing_quantiles = slice_rows(self.data["firing_quantiles"], unit_indices) new_bins = self.data["bins"][:] new_data = dict(acgs_3d=new_acgs_3d, firing_quantiles=new_firing_quantiles, bins=new_bins) return new_data @@ -1241,10 +1242,10 @@ def _merge_extension_data( new_firing_quantiles = np.zeros((len(new_sorting.unit_ids), firing_rate_quantiles.shape[1])) new_acgs_3d[new_unit_ids_indices, :, :] = acgs_3d - new_acgs_3d[old_unit_ids_indices, :, :] = self.data["acgs_3d"][old_unit_ids_indices, :, :] + new_acgs_3d[old_unit_ids_indices, :, :] = slice_rows(self.data["acgs_3d"], old_unit_ids_indices) new_firing_quantiles[new_unit_ids_indices, :] = firing_rate_quantiles - new_firing_quantiles[old_unit_ids_indices, :] = self.data["firing_quantiles"][old_unit_ids_indices, :] + new_firing_quantiles[old_unit_ids_indices, :] = slice_rows(self.data["firing_quantiles"], old_unit_ids_indices) new_data = dict( acgs_3d=new_acgs_3d, diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index ad4bc17594..f0a02ff7ec 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -10,7 +10,7 @@ import numpy as np from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension -from spikeinterface.core.core_tools import slice_rows +from spikeinterface.core.core_tools import slice_rows, materialize_array from spikeinterface.core.job_tools import TimeSeriesChunkExecutor, _shared_job_kwargs_doc, fix_job_kwargs from spikeinterface.core.analyzer_extension_core import _inplace_sparse_realign_waveforms @@ -91,7 +91,7 @@ def _select_units_extension_data(self, unit_ids): keep_spike_mask = np.isin(some_spikes["unit_index"], keep_unit_indices) new_data = dict() - new_data["pca_projection"] = self.data["pca_projection"][keep_spike_mask, :, :] + new_data["pca_projection"] = slice_rows(self.data["pca_projection"], keep_spike_mask) # one or several model for k, v in self.data.items(): if "model" in k: @@ -109,12 +109,15 @@ def _merge_extension_data( spike_indices = self.sorting_analyzer.get_extension("random_spikes").get_data() valid = keep_mask[spike_indices] some_spikes = some_spikes[valid] - pca_projections = pca_projections[valid] - else: - pca_projections = pca_projections.copy() + # slice_rows already returns an independent, materialized array + pca_projections = slice_rows(pca_projections, valid) old_sparsity = self.sorting_analyzer.sparsity if old_sparsity is not None: + if keep_mask is None: + # about to mutate pca_projections in place below (sparse realignment): we need a + # genuinely independent, writable buffer rather than sharing the original reference + pca_projections = materialize_array(pca_projections) # we need a realignement inside each group because we take the channel intersection sparsity # the story is same as in "waveforms" extension diff --git a/src/spikeinterface/postprocessing/tests/conftest.py b/src/spikeinterface/postprocessing/tests/conftest.py index 51ac8aa250..69fd5373bb 100644 --- a/src/spikeinterface/postprocessing/tests/conftest.py +++ b/src/spikeinterface/postprocessing/tests/conftest.py @@ -17,7 +17,7 @@ def _small_sorting_analyzer(): extensions_to_compute = { "random_spikes": {"seed": 1205}, - "noise_levels": {"seed": 1205}, + "noise_levels": {"random_slices_kwargs": {"seed": 1205}}, "waveforms": {}, "templates": {"operators": ["average", "median"]}, "spike_amplitudes": {}, diff --git a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py index cc2aabaefc..468e205772 100644 --- a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py +++ b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py @@ -5,15 +5,17 @@ from spikeinterface import ( create_sorting_analyzer, + load_sorting_analyzer, generate_ground_truth_recording, set_global_job_kwargs, get_template_amplitude_on_main_channel, ) from spikeinterface.core.generate import inject_some_split_units +from spikeinterface.core.core_tools import slice_rows # even if this is in postprocessing, we make an extension for quality metrics extension_dict = { - "noise_levels": dict(), + "noise_levels": dict(force_recompute=True), "random_spikes": dict(), "waveforms": dict(), "templates": dict(), @@ -59,7 +61,9 @@ "template_metrics": 0.2, # some metrics are very sensitive to template changes, so we put a large tolerance "quality_metrics": 1e-2, } -extensions_with_rel_tolerance_splits = {"amplitude_scalings": 1e-1} +extensions_with_rel_tolerance_splits = { + "amplitude_scalings": 1e-1, +} def get_dataset_to_merge(): @@ -122,8 +126,13 @@ def dataset_to_split(): return get_dataset_to_split() +@pytest.mark.parametrize("lazy", [False, True]) @pytest.mark.parametrize("sparse", [False, True]) -def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, sparse): +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, lazy, sparse, format, tmp_path): + if format == "memory" and lazy: + pytest.skip("lazy has no effect for format='memory' (nothing on disk to load lazily)") + set_global_job_kwargs(n_jobs=1) recording, sorting, other_ids = dataset_to_merge @@ -138,6 +147,13 @@ def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, sparse): sorting_analyzer.compute(extension_dict_merge, n_jobs=1) + if format != "memory": + analyzer_folder_name = f"sorting_analyzer_{sparse}_{lazy}" + if format == "zarr": + analyzer_folder_name += ".zarr" + sorting_analyzer.save_as(folder=tmp_path / analyzer_folder_name, format=format) + sorting_analyzer = load_sorting_analyzer(tmp_path / analyzer_folder_name, format=format, lazy=lazy) + # TODO: still some UserWarnings for n_jobs, where from? t0 = time.perf_counter() analyzer_merged_hard = sorting_analyzer.merge_units( @@ -183,7 +199,22 @@ def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, sparse): np.testing.assert_array_equal(data_original_unmerged, data_soft_unmerged) if ext not in random_computation: - np.testing.assert_array_equal(data_original_unmerged, data_hard_unmerged) + # unmerged units should be unchanged by a hard recompute; allow a tiny tolerance for + # floating point summation-order noise (e.g. different chunking/parallelization), + # not to be confused with a real discrepancy + if extension_data_type[ext] == "pandas": + original_for_hard_check = data_original_unmerged.dropna().to_numpy().astype("float") + hard_for_hard_check = data_hard_unmerged.dropna().to_numpy().astype("float") + else: + original_for_hard_check = data_original_unmerged + hard_for_hard_check = data_hard_unmerged + if original_for_hard_check.dtype.kind in ["U", "S", "O"]: + assert np.array_equal(original_for_hard_check, hard_for_hard_check) + elif original_for_hard_check.dtype.fields is None: + np.testing.assert_allclose(original_for_hard_check, hard_for_hard_check, rtol=1e-8, atol=1e-8) + else: + for f in original_for_hard_check.dtype.fields: + np.testing.assert_allclose(original_for_hard_check[f], hard_for_hard_check[f], rtol=1e-8, atol=1e-8) else: print(f"Skipping hard test for {ext} due to randomness in computation") @@ -219,16 +250,27 @@ def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, sparse): raise Exception(f"Failed for {ext} - field {f} - max error {max_error}") +@pytest.mark.parametrize("lazy", [False, True]) @pytest.mark.parametrize("sparse", [False, True]) -def test_SortingAnalyzer_split_all_extensions(dataset_to_split, sparse): +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +def test_SortingAnalyzer_split_all_extensions(dataset_to_split, lazy, sparse, format, tmp_path): + if format == "memory" and lazy: + pytest.skip("lazy has no effect for format='memory' (nothing on disk to load lazily)") set_global_job_kwargs(n_jobs=1) recording, sorting, units_to_split = dataset_to_split - sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=sparse) + sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=sparse, lazy=lazy) extension_dict_split = extension_dict.copy() sorting_analyzer.compute(extension_dict, n_jobs=1) + if format != "memory": + analyzer_folder_name = f"sorting_analyzer_{sparse}_{lazy}" + if format == "zarr": + analyzer_folder_name += ".zarr" + sorting_analyzer.save_as(folder=tmp_path / analyzer_folder_name, format=format) + sorting_analyzer = load_sorting_analyzer(tmp_path / analyzer_folder_name, format=format, lazy=lazy) + # we randomly apply splits (at half of spiketrain) num_spikes = sorting.count_num_spikes_per_unit() @@ -246,6 +288,9 @@ def test_SortingAnalyzer_split_all_extensions(dataset_to_split, sparse): extension_dict_ = extension_dict_split.copy() extension_dict_.pop("random_spikes") analyzer_hard.extensions["random_spikes"] = analyzer_split.extensions["random_spikes"] + # noise_levels' random slice sampling is seeded: reuse the exact same seed so a fresh + # recompute matches the original instead of sampling a different (but similarly valid) subset + extension_dict_["noise_levels"] = dict(sorting_analyzer.get_extension("noise_levels").params) analyzer_hard.compute(extension_dict_, n_jobs=1) for ext in extension_dict: @@ -310,14 +355,14 @@ def get_extension_data_for_units(sorting_analyzer, data, unit_ids, ext_data_type elif ext_data_type == "random": random_indices = sorting_analyzer.get_extension("random_spikes").get_data() unit_mask = np.isin(spike_vector[random_indices]["unit_index"], unit_indices) - return data[unit_mask] + return slice_rows(data, unit_mask) elif ext_data_type == "matrix": - return data[unit_indices][:, unit_indices] + return slice_rows(data, unit_indices)[:, unit_indices] elif ext_data_type == "unit": - return data[unit_indices] + return slice_rows(data, unit_indices) elif ext_data_type == "spike": unit_mask = np.isin(spike_vector["unit_index"], unit_indices) - return data[unit_mask] + return slice_rows(data, unit_mask) elif ext_data_type == "pandas": return data.loc[unit_ids].dropna()