Merge pull request #180 from SpikeInterface/oe-channel-subset

Open Ephys: Support subselection of channels in Record Node

Author: alejoe91

probeinterface/__init__.py

@@ -11,8 +11,9 @@
     read_imro, write_imro,
     read_BIDS_probe, write_BIDS_probe,
     read_spikeglx, parse_spikeglx_meta, get_saved_channel_indices_from_spikeglx_meta,
+    read_openephys, get_saved_channel_indices_from_openephys_settings,
     read_mearec, read_nwb,
-    read_maxwell, read_3brain, read_openephys)
+    read_maxwell, read_3brain)
 from .utils import combine_probes
 from .generator import (generate_dummy_probe, generate_dummy_probe_group,
             generate_tetrode, generate_linear_probe,

probeinterface/io.py

@@ -12,6 +12,7 @@
 from pathlib import Path
 from typing import Union, Optional
 import re
+import warnings
 import json
 from collections import OrderedDict
 from packaging.version import Version, parse
@@ -1262,12 +1263,14 @@ def read_openephys(
     oe_version = parse(info_chain.find("VERSION").text)
     signal_chain = root.find("SIGNALCHAIN")
     neuropix_pxi = None
+    record_node = None
     for processor in signal_chain:
         if "PROCESSOR" == processor.tag:
             name = processor.attrib["name"]
             if "Neuropix-PXI" in name:
                 neuropix_pxi = processor
-                break
+            if "Record Node" in name:
+                record_node = processor
 
     if neuropix_pxi is None:
         if raise_error:
@@ -1468,8 +1471,10 @@ def read_openephys(
     else:
         # in case of a single probe, make sure it is consistent with optional
         # stream_name, probe_name, or serial number
+        available_probe_name = np_probes_info[0]['name']
+        available_serial_number = np_probes_info[0]['serial_number']
+
         if stream_name:
-            available_probe_name = np_probes_info[0]['name']
             if available_probe_name not in stream_name:
                 if raise_error:
                     raise Exception(
@@ -1478,7 +1483,6 @@ def read_openephys(
                     )
                 return None
         if probe_name:
-            available_probe_name = np_probes_info[0]['name']
             if probe_name != available_probe_name:
                 if raise_error:
                     raise Exception(
@@ -1487,7 +1491,6 @@ def read_openephys(
                     )
                 return None
         if serial_number:
-            available_serial_number = np_probes_info[0]['serial_number']
             if str(serial_number) != available_serial_number:
                 if raise_error:
                     raise Exception(
@@ -1511,6 +1514,19 @@ def read_openephys(
         contact_width = 12
         shank_pitch = 250
 
+    contact_ids = np_probe_info['contact_ids'] if np_probe_info['contact_ids'] is not None else None
+
+    # check if subset of channels
+    chans_saved = get_saved_channel_indices_from_openephys_settings(settings_file, stream_name=stream_name)
+
+    # if a recording state is found, slice probe
+    if chans_saved is not None:
+        positions = positions[chans_saved]
+        if shank_ids is not None:
+            shank_ids = np.array(shank_ids)[chans_saved]
+        if contact_ids is not None:
+            contact_ids = np.array(contact_ids)[chans_saved]
+
     probe = Probe(ndim=2, si_units="um")
     probe.set_contacts(
         positions=positions,
@@ -1526,8 +1542,8 @@ def read_openephys(
         probe_serial_number=np_probe.attrib["probe_serial_number"],
     )
 
-    if np_probe_info['contact_ids'] is not None:
-        probe.set_contact_ids(np_probe_info['contact_ids'])
+    if contact_ids is not None:
+        probe.set_contact_ids(contact_ids)
 
     polygon = polygon_description["default"]
     if shank_ids is None:
@@ -1547,6 +1563,69 @@ def read_openephys(
     return probe
 
 
+def get_saved_channel_indices_from_openephys_settings(settings_file, stream_name):
+    """
+    Returns an array with the subset of saved channels indices (if used)
+
+    Parameters
+    ----------
+    settings_file : str or Path
+        The path to the settings file
+    stream_name : str
+        The stream name that contains the probe name
+        For example, "Record Node 100#ProbeA-AP" will select the AP stream of ProbeA.
+
+    Returns
+    -------
+    chans_saved
+        np.array of saved channel indices or None
+    """
+    # check if subset of channels
+    ET = import_safely("xml.etree.ElementTree")
+    # parse xml
+    tree = ET.parse(str(settings_file))
+    root = tree.getroot()
+
+    signal_chain = root.find("SIGNALCHAIN")
+    record_node = None
+    for processor in signal_chain:
+        if "PROCESSOR" == processor.tag:
+            name = processor.attrib["name"]
+            if "Record Node" in name:
+                record_node = processor
+                break
+    chans_saved = None
+    if record_node is not None:
+        custom_params = record_node.find("CUSTOM_PARAMETERS")
+        if custom_params is not None:
+            custom_streams = custom_params.findall("STREAM")
+            custom_stream_names = [stream.attrib["name"] for stream in custom_streams]
+            if len(custom_streams) > 0:
+                recording_states = [stream.attrib.get("recording_state", None) for stream in custom_streams]
+                has_custom_states = False
+                for rs in recording_states:
+                    if rs is not None and rs != "ALL":
+                        has_custom_states = True
+                if has_custom_states:
+                    if len(custom_streams) > 1:
+                        assert stream_name is not None, \
+                                (f"More than one stream found with custom parameters: {custom_stream_names}. "
+                                f"Use the `stream_name` argument to choose the correct stream")
+                        possible_custom_streams = [stream for stream in custom_streams
+                                                   if stream.attrib["name"] in stream_name]
+                        if len(possible_custom_streams) > 1:
+                            warnings.warn(f"More than one custom parameters associated to {stream_name} "
+                                          f"found. Using fisrt one")
+                        custom_stream = possible_custom_streams[0]
+                    else:
+                        custom_stream = custom_streams[0]
+                    recording_state = custom_stream.attrib.get("recording_state", None)
+                    if recording_state is not None:
+                        if recording_state != "ALL":
+                            chans_saved = np.array([chan for chan, r in enumerate(recording_state) if int(r) == 1])
+    return chans_saved
+
+
 def read_mearec(file: Union[str, Path]) -> Probe:
     """
     Read probe position, and contact shape from a MEArec file.