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..295f49fa31 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 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 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 : ZarrArrayExtractor + 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=ZarrArrayExtractor, name="read_zarr_array") def read_zarr(