Add Neuropixels 2.0 support to read_spikegadgets_neuropixels

src/probeinterface/io.py

@@ -733,17 +733,80 @@ def write_csv(file, probe):
     raise NotImplementedError
 
 
+def _spikegadgets_channel_index_np2_4shank(channel_index: int) -> int:
+    """Remap NP2.0 4-shank ``channelsOn`` bit position to catalogue index.
+
+    Trodes writes ``channelsOn`` row-major across all four shanks (eight
+    contacts per row, two columns per shank), with the column-within-row
+    direction reversed relative to ``probeColumn`` (high ``channel_index`` -> low
+    ``probeColumn``; see Trodes `configuration.cpp:5374-5421` and the
+    `probeColumn` annotations in the .rec ``SpikeChannel`` elements). The
+    catalogue (``NP2014``) is shank-major instead (s0e0..s0e1279,
+    s1e0..s1e1279, ...), so channel_index needs remapping. Verified empirically
+    against the SpikeGadgets-provided NP2.0 4-shank fixture: channel_index 1671 with
+    ``probeColumn="0"`` maps to ``s0e416``, channel_index 1664 with ``probeColumn="7"``
+    maps to ``s3e417``, and the .rec ``coord_ml``/``coord_dv`` values for those
+    SpikeChannel entries match the catalogue positions up to a single stereotactic
+    offset (these XML coords are not consumed by the reader, only used by the
+    test in `tests/test_io/test_spikegadgets.py` as an independent cross-check).
+    Independently confirmed in May 2026 by Mattias Karlsson (SpikeGadgets /
+    Trodes author) on PR #441: "the 2.0 four-shank probe has two columns per
+    shank... the first electrode on the probe (starts with 1) is in the lower
+    right, and number 10008 is on the lower left. Then, 10009 is the second
+    row on the right, and so on", which is exactly the (row, col_global=7-x)
+    layout this function encodes.
+    """
+    CONTACTS_PER_ROW = 8  # 2 columns per shank * 4 shanks
+    COLS_PER_SHANK = 2
+    CONTACTS_PER_SHANK = 1280
+
+    row = channel_index // CONTACTS_PER_ROW
+    col_global = (CONTACTS_PER_ROW - 1) - (channel_index % CONTACTS_PER_ROW)
+    shank = col_global // COLS_PER_SHANK
+    col_on_shank = col_global % COLS_PER_SHANK
+    return shank * CONTACTS_PER_SHANK + row * COLS_PER_SHANK + col_on_shank
+
+
+def _spikegadgets_channel_index_np2_1shank(channel_index: int) -> int:
+    """Remap NP2.0 single-shank ``channelsOn`` bit position to catalogue index.
+
+    Same row-major-within-probe layout as NP2.0 4-shank (Trodes
+    `configuration.cpp:5279-5290`) but with only one shank and two
+    columns per row, so two contacts per row. The within-row direction is
+    reversed relative to the catalogue (extrapolated from NP2.0 4-shank
+    where this was empirically verified): channel_index 0 -> right column, channel_index 1
+    -> left column, channel_index 2 -> next row right, etc. The catalogue
+    (``NP2000``) lays out contacts with left column first (idx 0 = left,
+    idx 1 = right per row), so the remap pairs are swapped:
+    catalogue_idx = row * 2 + (1 - channel_index % 2).
+
+    Unverified against a real fixture; will be revisited when a NP2.0
+    single-shank .rec from a Bennu rig becomes available.
+    """
+    COLS_PER_SHANK = 2
+
+    row = channel_index // COLS_PER_SHANK
+    col_on_shank = (COLS_PER_SHANK - 1) - (channel_index % COLS_PER_SHANK)
+    return row * COLS_PER_SHANK + col_on_shank
+
+
 def read_spikegadgets_neuropixels(file: str | Path, raise_error: bool = True) -> ProbeGroup:
     """
     Find active channels of the given Neuropixels probe from a SpikeGadgets .rec file.
-    SpikeGadgets headstages support up to three Neuropixels 1.0 probes (as of March 28, 2024),
+    SpikeGadgets headstages support up to three Neuropixels probes simultaneously,
     and information for all probes will be returned in a ProbeGroup object.
 
-    This function only supports Neuropixels probes recorded with SpikeGadgets
-    headstages (``HardwareConfiguration`` entries with ``name == "NeuroPixels1"``).
-    It does not handle tetrodes or other probe types that SpikeGadgets can
-    record. Use :func:`has_spikegadgets_neuropixels_probes` to check whether a
-    ``.rec`` file contains Neuropixels probe geometry before calling this reader.
+    Supported Neuropixels variants: NP1.0 standard (``device="neuropixels1"``,
+    older recordings without ``deviceSubType`` are treated as standard),
+    NP2.0 single-shank (``device="neuropixels2" deviceSubType="1_SHANK"``),
+    and NP2.0 4-shank (``device="neuropixels2" deviceSubType="4_SHANK"``).
+    The single-shank channel_index remap is extrapolated from the 4-shank pattern and
+    has not been verified against a real fixture yet. Other Neuropixels variants
+    Trodes can describe (NP1.0 HD, NP1.0 NHP short/medium/long, NRIC) raise
+    ``ValueError`` for now; non-Neuropixels probes (tetrodes etc.) are not
+    handled at all. Use :func:`has_spikegadgets_neuropixels_probes` to check
+    whether a ``.rec`` file contains Neuropixels probe geometry before calling
+    this reader.
 
     Parameters
     ----------
@@ -755,69 +818,149 @@ def read_spikegadgets_neuropixels(file: str | Path, raise_error: bool = True) ->
     probe_group : ProbeGroup object
 
     """
-    # The SpikeGadgets .rec XML does not include a probe part number. The NP1.0
-    # catalogue variants (NP1000, NP1001, PRB_1_2_0480_2, PRB_1_4_0480_1,
-    # PRB_1_4_0480_1_C) share identical 2D geometry in the probeinterface
-    # catalogue (contact positions, pitch, stagger, shank width), differing only
-    # in metadata that probeinterface does not consume (ADC resolution, databus
-    # phase, gain, on-shank reference, shank thickness). So hardcoding NP1000
-    # produces correct geometry; `model_name` and `description` are cleared on
-    # the sliced probe to avoid claiming a specific variant.
-    PART_NUMBER = "NP1000"
+    # Dispatch keyed by SpikeConfiguration (device, deviceSubType) attributes
+    # (see Trodes `configuration.cpp:2495-2520` and `5246-5291`). Each entry
+    # gives the HardwareConfiguration `Device` name to filter on, the catalogue
+    # part number to build the full probe from, the per-probe horizontal shift
+    # (um) used when plotting multi-probe ProbeGroups, and (optionally) a
+    # function remapping Trodes' ``channelsOn`` bit position (channel_index, equal to
+    # ``electrode_id[1:] - 1`` in the .rec XML) to a probeinterface catalogue
+    # contact index. The remap is None when Trodes' ordering already matches
+    # the catalogue's (NP1.0 standard).
+    #
+    # All NP1.0 staggered catalogue variants (NP1000, NP1001, NP1010-NP1014,
+    # PRB_1_2_0480_2, PRB_1_4_0480_1, PRB_1_4_0480_1_C) share identical 2D
+    # geometry, so NP1000 is the canonical pick. All NP2.0 4-shank catalogue
+    # variants (NP2010, NP2013, NP2014, NP2020, NP2021) share identical 2D
+    # geometry, so NP2014 is the canonical pick. model_name and description
+    # are cleared on the sliced probe in both cases because the XML does not
+    # carry a part-number field.
+    spikegadgets_neuropixels_formats = {
+        ("neuropixels1", "10"): {
+            "hardware_device_name": "NeuroPixels1",
+            "part_number": "NP1000",
+            "multi_probe_plot_offset_um": 250.0,
+            "channel_index_to_catalogue_index": None,
+        },
+        ("neuropixels2", "1_SHANK"): {
+            "hardware_device_name": "NeuroPixels2",
+            "part_number": "NP2000",
+            "multi_probe_plot_offset_um": 250.0,
+            "channel_index_to_catalogue_index": _spikegadgets_channel_index_np2_1shank,
+        },
+        ("neuropixels2", "4_SHANK"): {
+            "hardware_device_name": "NeuroPixels2",
+            "part_number": "NP2014",
+            "multi_probe_plot_offset_um": 1000.0,
+            "channel_index_to_catalogue_index": _spikegadgets_channel_index_np2_4shank,
+        },
+    }
 
     header_txt = parse_spikegadgets_header(file)
     root = ElementTree.fromstring(header_txt)
     hconf = root.find("HardwareConfiguration")
     sconf = root.find("SpikeConfiguration")
 
-    probe_configs = [d for d in hconf if d.attrib.get("name") == "NeuroPixels1"]
+    # Older NP1.0 recordings predate the device/deviceSubType attributes, so
+    # missing values fall back to NP1.0 standard.
+    sconf_device = (sconf.attrib.get("device", "") if sconf is not None else "").lower() or "neuropixels1"
+    sconf_subtype = sconf.attrib.get("deviceSubType", "") if sconf is not None else ""
+    if sconf_device == "neuropixels1" and not sconf_subtype:
+        sconf_subtype = "10"
+    dispatch_key = (sconf_device, sconf_subtype)
+    if dispatch_key not in spikegadgets_neuropixels_formats:
+        raise ValueError(
+            f"Unsupported SpikeGadgets Neuropixels variant device={sconf_device!r} "
+            f"deviceSubType={sconf_subtype!r}; supported: "
+            f"{sorted(spikegadgets_neuropixels_formats)}"
+        )
+    fmt = spikegadgets_neuropixels_formats[dispatch_key]
+
+    probe_configs = [d for d in hconf if d.attrib.get("name") == fmt["hardware_device_name"]]
     n_probes = len(probe_configs)
 
     if n_probes == 0:
         if raise_error:
-            raise Exception("No Neuropixels 1.0 probes found")
+            raise Exception(f"No {fmt['hardware_device_name']} probes found")
         return None
 
-    # NeuroPixels1 SourceOptions blocks carry the per-probe AP/LF gain settings.
-    # They appear in the same order as the SpikeNTrode probe digits (1, 2, 3).
-    source_options_blocks = [s for s in hconf.findall("SourceOptions") if s.attrib.get("name") == "NeuroPixels1"]
+    # SourceOptions blocks carry the per-probe AP/LF gain settings. They appear
+    # in the same order as the SpikeNTrode probe digits (1, 2, 3).
+    source_options_blocks = [
+        s for s in hconf.findall("SourceOptions") if s.attrib.get("name") == fmt["hardware_device_name"]
+    ]
 
     probe_group = ProbeGroup()
 
     for curr_probe in range(1, n_probes + 1):
         # SpikeNTrode elements are the authoritative list of recorded electrodes.
-        # Each id is "<probe_digit><1-based electrode number>" for up to 960
-        # electrodes on NP1.0; the catalogue uses 0-based indices, so
-        # catalogue_index = electrode_number - 1. The probe number is assumed
-        # to be a single digit (1, 2, or 3), matching the documented
-        # SpikeGadgets limit of three simultaneous Neuropixels probes.
+        # Each id is "<probe_digit><1-based electrode number>"; the leading digit
+        # identifies the probe (1, 2, or 3, matching the documented SpikeGadgets
+        # limit of three simultaneous Neuropixels probes) and the remainder is
+        # the 1-based electrode number on that probe (channel_index = electrode - 1).
+        # NP1.0 standard uses maxPadsPerProbe = 1000 (ids are 4 chars wide, e.g.
+        # "1384"); NP2.0 uses maxPadsPerProbe = 10000 (ids are 5 chars wide, e.g.
+        # "11672"). Slicing by [1:] handles both because the probe digit is
+        # always one char. The format's channel_index_to_catalogue_index function
+        # then remaps Trodes' channelsOn bit position to the catalogue's contact
+        # order; it is None when no remap is needed (NP1.0, where the catalogue
+        # happens to be in Trodes' bit order already).
         electrode_to_hwchan = {}
+        electrode_to_stereotactic = {}
         for ntrode in sconf:
             electrode_id = ntrode.attrib["id"]
             if int(electrode_id[0]) == curr_probe:
-                catalogue_index = int(electrode_id[1:]) - 1
-                hw_chan = int(ntrode[0].attrib["hwChan"])
-                electrode_to_hwchan[catalogue_index] = hw_chan
+                channel_index = int(electrode_id[1:]) - 1
+                catalogue_index = (
+                    channel_index
+                    if fmt["channel_index_to_catalogue_index"] is None
+                    else fmt["channel_index_to_catalogue_index"](channel_index)
+                )
+                spike_channel = ntrode[0]
+                electrode_to_hwchan[catalogue_index] = int(spike_channel.attrib["hwChan"])
+                electrode_to_stereotactic[catalogue_index] = (
+                    float(spike_channel.attrib["coord_ml"]),
+                    float(spike_channel.attrib["coord_dv"]),
+                    float(spike_channel.attrib["coord_ap"]),
+                )
 
         active_indices = np.array(sorted(electrode_to_hwchan.keys()))
 
-        full_probe = build_neuropixels_probe(PART_NUMBER)
+        full_probe = build_neuropixels_probe(fmt["part_number"])
         probe = full_probe.get_slice(active_indices)
 
-        # Clear part-number-specific metadata since we don't know the actual part number.
+        # Clear part-number-specific metadata since the .rec XML does not carry
+        # a part number; the catalogue pick is a geometry-equivalence stand-in
+        # rather than a fact read from the file.
         probe.model_name = ""
         probe.description = ""
+        probe.annotations.pop("part_number", None)
 
         device_channels = np.array([electrode_to_hwchan[idx] for idx in active_indices])
         probe.set_device_channel_indices(device_channels)
 
+        # Stereotactic coordinates from the .rec SpikeChannel attributes
+        # (workspace probe origin + on-probe offset, in micrometres; see Trodes
+        # `configuration.cpp:5443-5445`). These are recording-specific surgical
+        # metadata, distinct from `contact_positions` which carries the pure
+        # on-probe catalogue geometry. We attach them as per-contact annotations
+        # so downstream code that wants stereotactic locations (e.g. histology
+        # registration) can read them without re-parsing the XML.
+        stereotactic = np.array([electrode_to_stereotactic[idx] for idx in active_indices])
+        probe.annotate_contacts(
+            stereotactic_ml=stereotactic[:, 0],
+            stereotactic_dv=stereotactic[:, 1],
+            stereotactic_ap=stereotactic[:, 2],
+        )
+
         # Per-contact ADC group and sample order from the catalogue MUX table plus
         # the hwChan mapping (which is the readout-channel index for each contact).
         adc_sampling_table = probe.annotations.get("adc_sampling_table")
-        _annotate_probe_with_adc_sampling_info(probe, adc_sampling_table)
+        if adc_sampling_table is not None:
+            _annotate_probe_with_adc_sampling_info(probe, adc_sampling_table)
 
-        # NP1.0 gain is programmable. Read APGainMode and LFPGainMode from the
-        # SourceOptions block matching this probe (blocks appear in probe order).
+        # Neuropixels gain is programmable. Read APGainMode and LFPGainMode from
+        # the SourceOptions block matching this probe (blocks appear in probe order).
         if "ap_gain" not in probe.annotations and curr_probe - 1 < len(source_options_blocks):
             custom_options = {
                 opt.attrib["name"]: opt.attrib["data"].strip()
@@ -832,7 +975,7 @@ def read_spikegadgets_neuropixels(file: str | Path, raise_error: bool = True) ->
                     probe.annotate(lf_gain=float(lf_gain_str))
 
         # Shift multiple probes so they don't overlap when plotted
-        probe.move([250 * (curr_probe - 1), 0])
+        probe.move([fmt["multi_probe_plot_offset_um"] * (curr_probe - 1), 0])
 
         probe_group.add_probe(probe)
 
@@ -865,10 +1008,10 @@ def has_spikegadgets_neuropixels_probes(file: str | Path) -> bool:
 
     Detection scans the ``HardwareConfiguration`` block of the ``.rec`` XML
     header for ``Device`` entries whose ``name`` attribute matches a known
-    Neuropixels source name (currently ``"NeuroPixels1"``). The presence of
-    any such entry is the ground-truth signal that the file contains
-    Neuropixels probe geometry, independent of what other hardware the
-    headstage is also streaming.
+    Neuropixels source name (``"NeuroPixels1"`` or ``"NeuroPixels2"``). The
+    presence of any such entry is the ground-truth signal that the file
+    contains Neuropixels probe geometry, independent of what other hardware
+    the headstage is also streaming.
 
     Intended use: callers that route heterogeneous SpikeGadgets recordings
     (mixing tetrodes, Neuropixels, etc.) can gate the call to
@@ -884,6 +1027,8 @@ def has_spikegadgets_neuropixels_probes(file: str | Path) -> bool:
     -------
     bool
     """
+    neuropixels_source_names = {"NeuroPixels1", "NeuroPixels2"}
+
     try:
         header_txt = parse_spikegadgets_header(file)
         root = ElementTree.fromstring(header_txt)
@@ -895,7 +1040,7 @@ def has_spikegadgets_neuropixels_probes(file: str | Path) -> bool:
         return False
 
     for device in hconf:
-        if device.attrib.get("name") == "NeuroPixels1":
+        if device.attrib.get("name") in neuropixels_source_names:
             return True
     return False