Add ZarrArrayExtractor to read "plain" Zarr arrays

doc/api.rst

@@ -92,6 +92,7 @@ spikeinterface.core
     .. autofunction:: select_segment_sorting
     .. autofunction:: read_binary
     .. autofunction:: read_zarr
+    .. autofunction:: read_zarr_array
     .. autofunction:: apply_merges_to_sorting
     .. autofunction:: spike_vector_to_spike_trains
     .. autofunction:: random_spikes_selection

src/spikeinterface/core/zarrextractors.py

@@ -246,7 +246,11 @@ def write_recording(
 class ZarrRecordingSegment(BaseRecordingSegment):
     def __init__(self, root, dataset_name, **time_kwargs):
         BaseRecordingSegment.__init__(self, **time_kwargs)
-        self._timeseries = root[dataset_name]
+        if dataset_name is None:
+            # In this case, root is a simple array
+            self._timeseries = root
+        else:
+            self._timeseries = root[dataset_name]
 
     def get_num_samples(self) -> int:
         """Returns the number of samples in this signal block
@@ -268,6 +272,75 @@ def get_traces(
         return traces
 
 
+class ZarrArrayRecording(BaseRecording):
+    """
+    Recording class for a plain Zarr array with shape num_samples x num_channels.
+    Mimics loading a binary array using BinaryRecording.
+
+    Parameters
+    ----------
+    file_path : str
+        Path to the directory where the zarr array is stored
+    sampling_frequency : float
+        The sampling frequency
+    gain_to_uV : float or array-like, default: None
+        The gain to apply to the traces
+    offset_to_uV : float or array-like, default: None
+        The offset to apply to the traces
+    is_filtered : bool or None, default: None
+        If True, the recording is assumed to be filtered. If None, is_filtered is not set.
+    storage_options : dict or None: None
+        Storage options passed to the `zarr.open` function
+
+    Returns
+    -------
+    recording : ZarrArrayRecording
+        The recording Extractor
+    """
+
+    def __init__(
+        self,
+        file_path: str | Path,
+        sampling_frequency: float,
+        gain_to_uV: float | np.ndarray | None = None,
+        offset_to_uV: float | np.ndarray | None = None,
+        is_filtered: bool | None = None,
+        storage_options: dict | None = None,
+    ):
+
+        folder_path, _ = resolve_zarr_path(file_path)
+        self._root = super_zarr_open(folder_path, mode="r", storage_options=storage_options)
+
+        dtype = self._root.dtype
+        num_channels = self._root.shape[1]
+        channel_ids = list(range(num_channels))
+
+        BaseRecording.__init__(self, sampling_frequency, channel_ids, dtype)
+
+        rec_segment = ZarrRecordingSegment(self._root, None, sampling_frequency=sampling_frequency)
+        self.add_recording_segment(rec_segment)
+
+        if is_filtered is not None:
+            self.annotate(is_filtered=is_filtered)
+
+        if gain_to_uV is not None:
+            self.set_channel_gains(gain_to_uV)
+
+        if offset_to_uV is not None:
+            self.set_channel_offsets(offset_to_uV)
+
+        self._kwargs = {
+            "file_path": str(Path(file_path).absolute()),
+            "sampling_frequency": sampling_frequency,
+            "num_channels": num_channels,
+            "dtype": dtype.str,
+            "channel_ids": channel_ids,
+            "gain_to_uV": gain_to_uV,
+            "offset_to_uV": offset_to_uV,
+            "is_filtered": is_filtered,
+        }
+
+
 class _ZarrSegmentIndex:
     """Lazy segment_index array derived from segment_slices stored in zarr."""
 
@@ -485,6 +558,7 @@ def write_sorting(sorting: BaseSorting, folder_path: str | Path, storage_options
 
 read_zarr_recording = define_function_from_class(source_class=ZarrRecordingExtractor, name="read_zarr_recording")
 read_zarr_sorting = define_function_from_class(source_class=ZarrSortingExtractor, name="read_zarr_sorting")
+read_zarr_array = define_function_from_class(source_class=ZarrArrayRecording, name="read_zarr_array")
 
 
 def read_zarr(

src/spikeinterface/extractors/extractor_classes.py

@@ -14,6 +14,8 @@
     read_npy_snippets,
 )
 
+from spikeinterface.core.zarrextractors import read_zarr_array
+
 # sorting/recording/event from neo
 from .neoextractors import *
 from .neoextractors import read_neuroscope
@@ -199,5 +201,6 @@
         "read_zarr",
         "read_neuroscope",  # convenience function for neuroscope
         "read_split_intan_files",  # convenience function for segmented intan files
+        "read_zarr_array",
     ]
 )

src/spikeinterface/extractors/tests/test_zarr_array_extractor.py

@@ -0,0 +1,66 @@
+import numpy as np
+import pytest
+import zarr
+
+from spikeinterface.core.zarrextractors import read_zarr_array
+
+
+@pytest.fixture
+def make_dummy_zarr_data(tmp_path):
+    """Create 2D Zarr array on disk."""
+    zarr_path = tmp_path / "test_traces.zarr"
+    num_samples = 1000
+    num_channels = 4
+    dtype = np.int16
+
+    # dummy data
+    original_data = (np.arange(num_samples * num_channels, dtype=dtype) % 500).reshape(num_samples, num_channels)
+
+    z = zarr.open(
+        store=str(zarr_path),
+        mode="w",
+        shape=original_data.shape,
+        chunks=(200, num_channels),
+        dtype=dtype,
+    )
+    z[:] = original_data
+
+    return zarr_path, original_data
+
+
+def test_zarr_array_extractor(make_dummy_zarr_data):
+    """
+    Make a zarr array then reads it using `read_zarr_array` and
+    checks basic properties.
+    """
+    zarr_path, original_data = make_dummy_zarr_data
+    sampling_frequency = 30000.0
+    gain = 0.195
+    offset = 10.0
+
+    rec = read_zarr_array(
+        file_path=zarr_path,
+        sampling_frequency=sampling_frequency,
+        gain_to_uV=gain,
+        offset_to_uV=offset,
+        is_filtered=False,
+    )
+
+    assert rec.get_num_channels() == original_data.shape[1]
+    assert rec.get_num_samples() == original_data.shape[0]
+    assert rec.get_sampling_frequency() == sampling_frequency
+    assert np.all(rec.get_channel_gains() == gain)
+    assert np.all(rec.get_channel_offsets() == offset)
+
+    # 3. Verify exact trace reading
+    traces_raw = rec.get_traces(return_scaled=False)
+    np.testing.assert_array_equal(traces_raw, original_data)
+
+    # 4. Verify channel and time slicing
+    subset = rec.get_traces(start_frame=50, end_frame=150, channel_ids=[0, 2], return_in_uV=False)
+    np.testing.assert_array_equal(subset, original_data[50:150, [0, 2]])
+
+    # 5. Verify scaling math
+    traces_scaled = rec.get_traces(start_frame=0, end_frame=10, channel_ids=[1], return_in_uV=True)
+    expected_scaled = original_data[0:10, [1]] * gain + offset
+    np.testing.assert_allclose(traces_scaled, expected_scaled, rtol=1e-5)