Propagate lazy arg to from si.load() (#4709)

Author: alejoe91

src/spikeinterface/core/core_tools.py

@@ -13,6 +13,7 @@
 
 from probeinterface import ProbeGroup
 import numpy as np
+import zarr
 
 
 def define_function_handling_dict_from_class(source_class, name):
@@ -779,3 +780,26 @@ def is_path_remote(path: str | Path) -> bool:
 def ms_to_samples(ms: float, sampling_frequency: float) -> int:
     """Convert a duration in milliseconds to the nearest number of samples."""
     return round(ms * sampling_frequency / 1000.0)
+
+
+def slice_rows(array: np.ndarray | zarr.Array, row_indices: np.ndarray | list) -> np.ndarray:
+    """
+    Slice a 2D array to select specific rows based on provided indices.
+
+    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.
+
+    Returns
+    -------
+    np.ndarray
+        A new 2D numpy array containing only the selected rows.
+    """
+    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]
+    else:
+        return array[row_indices, ...]

src/spikeinterface/core/loading.py

@@ -282,11 +282,7 @@ def _load_object_from_zarr(folder_or_url, object_type, **kwargs):
     if object_type == "SortingAnalyzer":
         from .sortinganalyzer import load_sorting_analyzer
 
-        backend_options = kwargs.get("backend_options", None)
-        load_extensions = kwargs.get("load_extensions", True)
-        analyzer = load_sorting_analyzer(
-            folder_or_url, backend_options=backend_options, load_extensions=load_extensions
-        )
+        analyzer = load_sorting_analyzer(folder_or_url, **kwargs)
         return analyzer
     elif object_type == "Templates":
         from .template import Templates

src/spikeinterface/core/sortinganalyzer.py

@@ -466,14 +466,21 @@ def __repr__(self) -> str:
         nchan = self.get_num_channels()
         nunits = self.get_num_units()
         txt = f"{clsname}: {nchan} channels - {nunits} units - {nseg} segments - {self.format}"
-        if self.format != "memory" and is_path_remote(self.folder):
-            txt += " (remote)"
+        if self.format != "memory":
+            if is_path_remote(self.folder):
+                if self._lazy:
+                    txt += " (remote + lazy)"
+                else:
+                    txt += " (remote)"
+            elif self._lazy:
+                txt += " (lazy)"
         if self.is_sparse():
             txt += " - sparse"
         if self.has_recording():
             txt += " - has recording"
         if self.has_temporary_recording():
             txt += " - has temporary recording"
+
         ext_txt = f"Loaded {len(self.extensions)} extensions"
         if len(self.extensions) > 0:
             ext_txt += f": {', '.join(self.extensions.keys())}"

src/spikeinterface/postprocessing/principal_component.py

@@ -10,9 +10,8 @@
 import numpy as np
 
 from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension
-
+from spikeinterface.core.core_tools import slice_rows
 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
 
 _possible_modes = ["by_channel_local", "by_channel_global", "concatenated"]
@@ -197,7 +196,7 @@ def get_projections_one_unit(self, unit_id, sparse=False):
 
         unit_index = sorting.id_to_index(unit_id)
         spike_mask = some_spikes["unit_index"] == unit_index
-        projections = self.data["pca_projection"][spike_mask]
+        projections = slice_rows(self.data["pca_projection"], spike_mask)
 
         if sparsity is None:
             return projections

src/spikeinterface/postprocessing/valid_unit_periods.py

@@ -9,6 +9,7 @@
 from tqdm.auto import tqdm
 
 from spikeinterface.core.base import unit_period_dtype
+from spikeinterface.core.core_tools import slice_rows
 from spikeinterface.core.job_tools import fix_job_kwargs
 from spikeinterface.core.sorting_tools import cast_periods_to_unit_period_dtype, remap_unit_indices_in_vector
 from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension
@@ -539,7 +540,7 @@ def _get_data(self, outputs: str = "by_unit"):
             (start_sample_index, end_sample_index) tuples.
         """
         if outputs == "numpy":
-            good_periods = self.data["valid_unit_periods"].copy()
+            good_periods = np.asarray(self.data["valid_unit_periods"]).copy()
         else:
             # by_unit
             unit_ids = self.sorting_analyzer.unit_ids
@@ -551,7 +552,7 @@ def _get_data(self, outputs: str = "by_unit"):
                 for unit_index, unit_id in enumerate(unit_ids):
                     periods_dict[unit_id] = []
                     unit_mask = good_periods_array["unit_index"] == unit_index
-                    good_periods_unit_segment = good_periods_array[segment_mask & unit_mask]
+                    good_periods_unit_segment = slice_rows(good_periods_array, segment_mask & unit_mask)
                     for start, end in good_periods_unit_segment[["start_sample_index", "end_sample_index"]]:
                         periods_dict[unit_id].append((start, end))
                 good_periods.append(periods_dict)