[DRAFT] Optimize quality metrics performance

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
 -----
 

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
 

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

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(