diff --git a/doc/api.rst b/doc/api.rst index ce850d1291..a51e26ac91 100755 --- a/doc/api.rst +++ b/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 diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 6cdc1c9fde..640d10e1c8 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/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( diff --git a/src/spikeinterface/extractors/extractor_classes.py b/src/spikeinterface/extractors/extractor_classes.py index 40df2716d2..1e9f5f6784 100644 --- a/src/spikeinterface/extractors/extractor_classes.py +++ b/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", ] ) diff --git a/src/spikeinterface/extractors/tests/test_zarr_array_extractor.py b/src/spikeinterface/extractors/tests/test_zarr_array_extractor.py new file mode 100644 index 0000000000..ddab87d4e1 --- /dev/null +++ b/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)