From e4e434bad692f17b2b5cff5e4fd920bb47c7f149 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 4 Sep 2026 09:59:26 +0100 Subject: [PATCH 1/8] add ZarrArrayExtractor --- doc/api.rst | 1 + src/spikeinterface/core/zarrextractors.py | 90 ++++++++++++++++++++++- 2 files changed, 90 insertions(+), 1 deletion(-) 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..3cba40711e 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,89 @@ def get_traces( return traces +class ZarrArrayExtractor(BaseRecording): + """ + RecordingExtractor for a plain Zarr array with shape num_samples x num_channels. + Mimics loading a binary array using BinaryRecordingExtractor. + + Parameters + ---------- + file_path : str + Path to the binary file + sampling_frequency : float + The sampling frequency + num_channels : int + Number of channels + dtype : str or dtype + The dtype of the binary file + channel_ids : list, default: None + A list of channel ids + 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 : ZarrArrayExtractor + The recording Extractor + """ + + def __init__( + self, + file_path, + sampling_frequency, + dtype, + num_channels: int | None = None, + channel_ids=None, + gain_to_uV=None, + offset_to_uV=None, + is_filtered=None, + storage_options=None, + ): + + assert num_channels is not None, "`num_channels` must be given." + + if channel_ids is None: + channel_ids = list(range(num_channels)) + else: + assert len(channel_ids) == num_channels, "Provided recording channels have the wrong length" + + BaseRecording.__init__(self, sampling_frequency, channel_ids, dtype) + + dtype = np.dtype(dtype) + + folder_path, _ = resolve_zarr_path(file_path) + self._root = super_zarr_open(folder_path, mode="r", storage_options=storage_options) + + 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 +572,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=ZarrArrayExtractor, name="read_zarr_array") def read_zarr( From 81dd4e73c8210fb8d433361a99f0e8e823479226 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 4 Sep 2026 10:07:39 +0100 Subject: [PATCH 2/8] remove custom channel_ids --- src/spikeinterface/core/zarrextractors.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 3cba40711e..ec88dea66e 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -287,8 +287,6 @@ class ZarrArrayExtractor(BaseRecording): Number of channels dtype : str or dtype The dtype of the binary file - channel_ids : list, default: None - A list of channel ids 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 @@ -318,11 +316,7 @@ def __init__( ): assert num_channels is not None, "`num_channels` must be given." - - if channel_ids is None: - channel_ids = list(range(num_channels)) - else: - assert len(channel_ids) == num_channels, "Provided recording channels have the wrong length" + channel_ids = list(range(num_channels)) BaseRecording.__init__(self, sampling_frequency, channel_ids, dtype) From 4293022fe5d5e03cfea85caa0f6e0dc36635922a Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 4 Sep 2026 12:16:40 +0100 Subject: [PATCH 3/8] remove dtype and num_channels --- src/spikeinterface/core/zarrextractors.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index ec88dea66e..2c5ebe0814 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -283,10 +283,6 @@ class ZarrArrayExtractor(BaseRecording): Path to the binary file sampling_frequency : float The sampling frequency - num_channels : int - Number of channels - dtype : str or dtype - The dtype of the binary file 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 @@ -306,25 +302,21 @@ def __init__( self, file_path, sampling_frequency, - dtype, - num_channels: int | None = None, - channel_ids=None, gain_to_uV=None, offset_to_uV=None, is_filtered=None, storage_options=None, ): - assert num_channels is not None, "`num_channels` must be given." + 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) - dtype = np.dtype(dtype) - - folder_path, _ = resolve_zarr_path(file_path) - self._root = super_zarr_open(folder_path, mode="r", storage_options=storage_options) - rec_segment = ZarrRecordingSegment(self._root, None, sampling_frequency=sampling_frequency) self.add_recording_segment(rec_segment) From c281eaef9a2c87b51554bef67ec28ef65b71b4b4 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 4 Sep 2026 12:21:37 +0100 Subject: [PATCH 4/8] typing and docstring --- src/spikeinterface/core/zarrextractors.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 2c5ebe0814..295f49fa31 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -280,7 +280,7 @@ class ZarrArrayExtractor(BaseRecording): Parameters ---------- file_path : str - Path to the binary file + 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 @@ -300,12 +300,12 @@ class ZarrArrayExtractor(BaseRecording): def __init__( self, - file_path, - sampling_frequency, - gain_to_uV=None, - offset_to_uV=None, - is_filtered=None, - storage_options=None, + 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) From 2e7b5ee513d3c3d9e1b291ca475e46b98840ee00 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Thu, 10 Sep 2026 07:27:12 +0100 Subject: [PATCH 5/8] add test --- .../tests/test_zarr_array_extractor.py | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 src/spikeinterface/extractors/tests/test_zarr_array_extractor.py 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) From f7863bb2019a1452f71dca48512e1cb7eef823cc Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Thu, 10 Sep 2026 07:27:20 +0100 Subject: [PATCH 6/8] add to init --- src/spikeinterface/extractors/extractor_classes.py | 3 +++ 1 file changed, 3 insertions(+) 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", ] ) From 489591f211b5245a974cac8a4ad5d286288031ad Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Thu, 10 Sep 2026 12:25:57 +0100 Subject: [PATCH 7/8] extractor -> recording --- src/spikeinterface/core/zarrextractors.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 295f49fa31..8de7b96a9d 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -272,10 +272,10 @@ def get_traces( return traces -class ZarrArrayExtractor(BaseRecording): +class ZarrArrayRecording(BaseRecording): """ - RecordingExtractor for a plain Zarr array with shape num_samples x num_channels. - Mimics loading a binary array using BinaryRecordingExtractor. + Recording class for a plain Zarr array with shape num_samples x num_channels. + Mimics loading a binary array using BinaryRecording. Parameters ---------- From 5c1e6a74dff314589063f36931d349d835a5ed5d Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Thu, 10 Sep 2026 12:31:12 +0100 Subject: [PATCH 8/8] oups --- src/spikeinterface/core/zarrextractors.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 8de7b96a9d..640d10e1c8 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -294,7 +294,7 @@ class ZarrArrayRecording(BaseRecording): Returns ------- - recording : ZarrArrayExtractor + recording : ZarrArrayRecording The recording Extractor """ @@ -558,7 +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=ZarrArrayExtractor, name="read_zarr_array") +read_zarr_array = define_function_from_class(source_class=ZarrArrayRecording, name="read_zarr_array") def read_zarr(