Add OpenEphysArrowRecording class

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),

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).