add test and move mux logic to read_spikeglx

Author: h-mayorquin

.gitignore

@@ -27,3 +27,4 @@ cov.xml
 
 
 neuropixel_library_generated/*
+uv.lock

src/probeinterface/neuropixels_tools.py

@@ -471,31 +471,16 @@ def build_neuropixels_probe(probe_part_number: str) -> Probe:
         lf_sample_frequency_hz=float(probe_spec_dict["lf_sample_frequency_hz"]),
     )
 
-    # ===== 8. Add MUX table annotations =====
-    mux_table_string = probe_features["z_mux_tables"][probe_spec_dict["mux_table_format_type"]]
-    if mux_table_string is not None:
-        # Parse MUX table string: (num_adcs,num_channels_per_adc)(int int ...)(int int ...)...
-        adc_info = mux_table_string.split(")(")[0]
-        split_mux = mux_table_string.split(")(")[1:]
-        num_adcs, num_channels_per_adc = map(int, adc_info[1:].split(","))
-        adc_groups_list = [
-            np.array(each_mux.replace("(", "").replace(")", "").split(" ")).astype("int") for each_mux in split_mux
-        ]
-        mux_table = np.transpose(np.array(adc_groups_list))
-
-        # Map contacts to ADC groups and sample order
-        num_contacts = positions.shape[0]
-        adc_groups = np.zeros(num_contacts, dtype="int64")
-        adc_sample_order = np.zeros(num_contacts, dtype="int64")
-        for adc_index, adc_groups_per_adc in enumerate(mux_table):
-            adc_groups_per_adc = adc_groups_per_adc[adc_groups_per_adc < num_contacts]
-            adc_groups[adc_groups_per_adc] = adc_index
-            adc_sample_order[adc_groups_per_adc] = np.arange(len(adc_groups_per_adc))
-
-        probe.annotate(num_adcs=num_adcs)
-        probe.annotate(num_channels_per_adc=num_channels_per_adc)
-        probe.annotate_contacts(adc_group=adc_groups)
-        probe.annotate_contacts(adc_sample_order=adc_sample_order)
+    # ===== 8. Store ADC sampling table =====
+    # The ADC sampling table describes how readout channels map to ADCs, not electrodes.
+    # Per-contact annotations (adc_group, adc_sample_order) can only be correctly
+    # assigned when reading a recording with a known channel map (via read_spikeglx),
+    # because the table indices are readout channel indices, not electrode indices.
+    # We store the full table string so it's available after slicing.
+    mux_table_format_type = probe_spec_dict["mux_table_format_type"]
+    adc_sampling_table = probe_features["z_mux_tables"].get(mux_table_format_type)
+    if adc_sampling_table is not None:
+        probe.annotate(adc_sampling_table=adc_sampling_table)
 
     return probe
 
@@ -861,6 +846,38 @@ def read_spikeglx(file: str | Path) -> Probe:
         annotations[pi_field] = values
     probe.annotate_contacts(**annotations)
 
+    # ===== 6b. Add ADC sampling annotations =====
+    # The ADC sampling table describes which ADC samples each readout channel and in what order.
+    # At this point, contacts are ordered by readout channel (0-383), so we can directly
+    # apply the mapping. This must be done here (not in build_neuropixels_probe)
+    # because the table indices are readout channel indices, not electrode indices.
+    adc_sampling_table = probe.annotations.get("adc_sampling_table")
+    if adc_sampling_table is not None:
+        # Parse table string: (num_adcs,num_channels_per_adc)(ch ch ...)(ch ch ...)...
+        adc_info = adc_sampling_table.split(")(")[0]
+        split_mux = adc_sampling_table.split(")(")[1:]
+        num_adcs, num_channels_per_adc = map(int, adc_info[1:].split(","))
+        adc_groups_list = [
+            np.array(each_mux.replace("(", "").replace(")", "").split(" ")).astype("int")
+            for each_mux in split_mux
+        ]
+        mux_table = np.transpose(np.array(adc_groups_list))
+
+        # Map readout channels to ADC groups and sample order
+        num_readout_channels = probe.get_contact_count()
+        adc_groups = np.zeros(num_readout_channels, dtype="int64")
+        adc_sample_order = np.zeros(num_readout_channels, dtype="int64")
+        for adc_index, channels_per_adc in enumerate(mux_table):
+            # Filter out placeholder values (e.g., 128 in mux_np1200 for unused slots)
+            valid_channels = channels_per_adc[channels_per_adc < num_readout_channels]
+            adc_groups[valid_channels] = adc_index
+            adc_sample_order[valid_channels] = np.arange(len(valid_channels))
+
+        probe.annotate(num_adcs=num_adcs)
+        probe.annotate(num_channels_per_adc=num_channels_per_adc)
+        probe.annotate_contacts(adc_group=adc_groups)
+        probe.annotate_contacts(adc_sample_order=adc_sample_order)
+
     # ===== 7. Slice to saved channels (if subset was saved) =====
     # This is DIFFERENT from IMRO selection: IMRO selects which electrodes to acquire,
     # but SpikeGLX can optionally save only a subset of acquired channels to reduce file size.

tests/test_io/test_spikeglx.py

@@ -320,6 +320,57 @@ def test_CatGT_NP1():
     assert "1.0" in probe.description
 
 
+def test_mux_annotations_match_readout_channels():
+    """
+    Test that MUX annotations (adc_group, adc_sample_order) correctly map
+    readout channels to their ADC assignments.
+
+    The MUX table defines which ADC samples each readout channel and in what order.
+    For NP1.0 (mux_np1000), the pattern is:
+    - Readout channel 0 -> ADC 0, sample order 0
+    - Readout channel 1 -> ADC 1, sample order 0
+    - Readout channel 2 -> ADC 0, sample order 1
+    - Readout channel 3 -> ADC 1, sample order 1
+    - etc.
+
+    The contacts in the probe returned by read_spikeglx are ordered by readout channel,
+    so adc_group[i] should be the ADC for readout channel i.
+
+    This test uses allan-longcol which has alternating banks (0, 1, 0, 1, ...),
+    meaning electrode indices differ from readout channel indices. This exposes
+    bugs where MUX table indices (readout channels) are incorrectly treated as
+    electrode indices.
+    """
+    # allan-longcol uses alternating banks: channel 0->electrode 0, channel 1->electrode 385, etc.
+    probe = read_spikeglx(data_path / "allan-longcol_g0_t0.imec0.ap.meta")
+
+    assert probe.get_contact_count() == 384
+
+    # Verify the bank pattern - this file alternates between bank 0 and bank 1
+    banks = probe.contact_annotations.get("banks")
+    assert list(banks[:10]) == [0, 1, 0, 1, 0, 1, 0, 1, 0, 1], "Expected alternating banks"
+
+    adc_groups = probe.contact_annotations.get("adc_group")
+    adc_sample_order = probe.contact_annotations.get("adc_sample_order")
+
+    assert adc_groups is not None, "adc_group annotation should be present"
+    assert adc_sample_order is not None, "adc_sample_order annotation should be present"
+    assert len(adc_groups) == 384
+    assert len(adc_sample_order) == 384
+
+    # According to mux_np1000 (32 ADCs, 12 channels per ADC):
+    # Readout channel 0 -> ADC 0, readout channel 1 -> ADC 1, etc.
+    # The MUX assignment follows readout channel order, NOT electrode order.
+    # Even though electrode indices are [0, 385, 2, 387, ...], the ADC groups
+    # should still be [0, 1, 0, 1, ...] because that's the readout channel pattern.
+    expected_adc_for_first_10 = [0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
+    assert list(adc_groups[:10]) == expected_adc_for_first_10, (
+        f"ADC groups for first 10 channels should be {expected_adc_for_first_10}, "
+        f"got {list(adc_groups[:10])}. "
+        f"MUX annotations must follow readout channel order, not electrode order."
+    )
+
+
 def test_snsGeomMap():
     # check when snsGeomMap is present if contact positions are the same from imroTbl