From e4e434bad692f17b2b5cff5e4fd920bb47c7f149 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 4 Sep 2026 09:59:26 +0100 Subject: [PATCH 1/4] 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/4] 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/4] 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/4] 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)