diff --git a/src/spikeinterface/extractors/neoextractors/__init__.py b/src/spikeinterface/extractors/neoextractors/__init__.py index a90e0954f3..6098d3e2a2 100644 --- a/src/spikeinterface/extractors/neoextractors/__init__.py +++ b/src/spikeinterface/extractors/neoextractors/__init__.py @@ -24,8 +24,11 @@ OpenEphysLegacyRecordingExtractor, OpenEphysBinaryRecordingExtractor, OpenEphysBinaryEventExtractor, + # we treat OpenEphysArrowRecording like a neo extractor + OpenEphysArrowRecording, read_openephys, read_openephys_event, + read_openephys_arrow, ) from .plexon import PlexonRecordingExtractor, PlexonSortingExtractor, read_plexon, read_plexon_sorting from .plexon2 import ( @@ -63,6 +66,7 @@ NixRecordingExtractor: dict(wrapper_string="read_nix", wrapper_class=read_nix), OpenEphysBinaryRecordingExtractor: dict(wrapper_string="read_openephys", wrapper_class=read_openephys), OpenEphysLegacyRecordingExtractor: dict(wrapper_string="read_openephys", wrapper_class=read_openephys), + OpenEphysArrowRecording: dict(wrapper_string="read_openephys_arrow", wrapper_class=read_openephys_arrow), PlexonRecordingExtractor: dict(wrapper_string="read_plexon", wrapper_class=read_plexon), Plexon2RecordingExtractor: dict(wrapper_string="read_plexon2", wrapper_class=read_plexon2), Spike2RecordingExtractor: dict(wrapper_string="read_spike2", wrapper_class=read_spike2), diff --git a/src/spikeinterface/extractors/neoextractors/openephys.py b/src/spikeinterface/extractors/neoextractors/openephys.py index 824f8d9560..de17909c90 100644 --- a/src/spikeinterface/extractors/neoextractors/openephys.py +++ b/src/spikeinterface/extractors/neoextractors/openephys.py @@ -8,6 +8,7 @@ for more info. """ +import importlib.util from pathlib import Path import numpy as np @@ -21,6 +22,9 @@ ) from spikeinterface.extractors.neoextractors.neobaseextractor import NeoBaseRecordingExtractor, NeoBaseEventExtractor +from spikeinterface.core.core_tools import define_function_from_class +from spikeinterface.core import BaseRecording, BaseRecordingSegment + def drop_invalid_neo_arguments_for_version_0_12_0(neo_kwargs): from packaging.version import Version @@ -491,6 +495,122 @@ def map_to_neo_kwargs(cls, folder_path, experiment_names=None): return neo_kwargs +class OpenEphysArrowRecordingSegment(BaseRecordingSegment): + def __init__(self, dataset, channel_ids, **time_kwargs): + BaseRecordingSegment.__init__(self, **time_kwargs) + self._dataset = dataset + self._all_channel_ids = channel_ids + + def get_num_samples(self) -> int: + """Returns the number of samples in this signal block + + Returns: + SampleIndex : Number of samples in the signal block + """ + return self._dataset.count_rows() + + def get_traces( + self, + start_frame: int | None = None, + end_frame: int | None = None, + channel_indices: list[int | str] | None = None, + ) -> np.ndarray: + if channel_indices is None: + channel_ids = list(self._all_channel_ids) + else: + channel_ids = list(self._all_channel_ids[channel_indices]) + + scanner = self._dataset.scanner(columns=channel_ids) + return np.column_stack(scanner.take(range(start_frame, end_frame))) + + +class OpenEphysArrowRecording(BaseRecording): + """ + Recording class for the openephys arrow format, from + + Parameters + ---------- + file_path : str + Path to the directory where the zarr array is stored + sampling_frequency : float + The sampling frequency + stream_name : str, default: AmplifierData + The stream name of the data you want to load. By default, the ephys AP stream is + called "AmplifierData". + 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, + stream_name="AmplifierData", + gain_to_uV: float | np.ndarray | None = None, + offset_to_uV: float | np.ndarray | None = None, + is_filtered: bool | None = None, + ): + if importlib.util.find_spec("pyarrow") is None: + raise ImportError("You need to add `pyarrow` to your environment to open .arrow files") + else: + import pyarrow.dataset as ds + + dataset = ds.dataset(file_path, format="arrow") + stream_names = dataset.schema.names + channel_ids = [name for name in stream_names if stream_name in name] + + if len(channel_ids) == 0: + raise ValueError(f"Cannot find any data with `stream_name` = {stream_name}") + + one_channel_index = stream_names.index(channel_ids[0]) + + # Arrow uses it's own DataType. For ints, it converts to numpy dtype without issue + ephys_type = dataset.schema[one_channel_index].type + numpy_type = np.dtype(str(ephys_type)) + + BaseRecording.__init__(self, sampling_frequency=sampling_frequency, channel_ids=channel_ids, dtype=numpy_type) + + rec_segment = OpenEphysArrowRecordingSegment( + dataset, sampling_frequency=sampling_frequency, channel_ids=np.array(channel_ids) + ) + + 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": len(channel_ids), + "dtype": numpy_type.str, + "channel_ids": channel_ids, + "gain_to_uV": gain_to_uV, + "offset_to_uV": offset_to_uV, + "is_filtered": is_filtered, + } + + +read_openephys_arrow = define_function_from_class(source_class=OpenEphysArrowRecording, name="read_openephys_arrow") + + def read_openephys(folder_path, **kwargs): """ Read Open Ephys folder (in "binary" or "open ephys legacy" format).