Add an 'array' field to the channel-map ch axis

src/ezmsg/blackrock/cerelink.py

@@ -18,7 +18,7 @@
 from ezmsg.util.messages.axisarray import AxisArray, replace
 from pycbsdk import ChanInfoField, ChannelType, DeviceType, SampleRate, Session
 
-from .channel_map import CHANNEL_DTYPE, ChannelMapSettings
+from .channel_map import CHANNEL_DTYPE, ChannelMapSettings, _array_identity
 from .clock import device_to_monotonic_batch_offsets
 
 logger = logging.getLogger(__name__)
@@ -413,6 +413,7 @@ def _build_ch_info(self, channels: list[int]) -> np.ndarray:
             ch_info[i]["bank"] = chr(ord("A") + bank_num - 1) if bank_num > 0 else ""
             ch_info[i]["elec"] = term
             ch_info[i]["headstage"] = headstage
+            ch_info[i]["array"] = _array_identity(ch_info[i]["label"], headstage, ch_info[i]["bank"])
         return ch_info
 
     def _device_name(self) -> str:

src/ezmsg/blackrock/channel_map.py

@@ -1,9 +1,12 @@
 """Attach Blackrock ``.cmp`` channel-map metadata to an ``AxisArray``'s ``ch`` axis.
 
 The output ``ch`` axis is a structured ``CoordinateAxis`` with fields
-``x``, ``y``, ``size``, ``label``, ``bank``, ``elec``, ``headstage`` for every
-input channel. ``x``/``y``/``size`` are in micrometers; ``headstage`` is the
-1-based headstage id (``0`` = none/auto).
+``x``, ``y``, ``size``, ``label``, ``bank``, ``elec``, ``headstage``, ``array``
+for every input channel. ``x``/``y``/``size`` are in micrometers;
+``headstage`` is the 1-based headstage id (``0`` = none/auto); ``array`` is the
+electrode-array identity derived from the label (see :func:`_array_identity`),
+which groups channels by physical array rather than by 32-channel connector
+bank — the useful grouping for per-array rereferencing.
 
 :class:`ChannelMapUnit` takes the *complete* set of per-headstage overlays in
 one settings object (:class:`ChannelMapUnitSettings`, a tuple of
@@ -64,10 +67,52 @@
         ("bank", "U1"),
         ("elec", "i4"),
         ("headstage", "i4"),  # 1-based headstage id (0 = none/auto)
+        # Electrode-array identity, e.g. "hs1-elec1-m1". One value per
+        # physically distinct array; see _array_identity.
+        ("array", "U32"),
     ]
 )
 
 
+def _array_identity(label: str, headstage: int, bank: str = "") -> str:
+    """Electrode-array identity for a channel.
+
+    Two channels belong to the same array iff they share a *connector* -- the
+    label prefix before the first ``-``, which the CMP assigns per physical
+    64-channel connector (``elec1`` in ``elec1-m1-63``) -- on the same
+    headstage. The connector label repeats across headstages, so the headstage
+    is prefixed when known; without it the same ``.cmp`` loaded twice, or a
+    bilateral implant, would merge two physically distinct arrays::
+
+        ("elec1-m1-63", 1) -> "hs1-elec1"
+        ("elec1-m1-63", 2) -> "hs2-elec1"
+        ("elec1-m1-63", 0) -> "elec1"
+
+    Grouping on the connector rather than the whole label stem is deliberate:
+    the region token is an annotation of where the array was implanted, and a
+    connector whose channels carry inconsistent region labels is still one
+    array. It also matches the offline convention in intent-pipelines
+    (``_array_ids_from_metadata``), so weights fitted offline and applied live
+    see the same clusters.
+
+    Labels with no connector structure (``chan1``, or an auto-grid channel with
+    no label) fall back to the connector *bank*::
+
+        ("chan1", 1, "A") -> "hs1-bankA"
+
+    so grouping by ``array`` degrades to bank-level grouping rather than
+    collapsing every such channel into one cluster spanning the whole device.
+    Returns ``""`` only when neither a label nor a bank is available.
+    """
+    hs = int(headstage or 0)
+    prefix = f"hs{hs}-" if hs > 0 else ""
+    connector, sep, _rest = str(label or "").partition("-")
+    if sep and connector:
+        return f"{prefix}{connector}"
+    bank = str(bank or "")
+    return f"{prefix}bank{bank}" if bank else ""
+
+
 class ChannelMapSettings(ez.Settings):
     filepath: str | None = None
     """Path to the ``.cmp`` file. ``None`` (or an empty path) means no CMP —
@@ -158,6 +203,7 @@ def _reset_state(self, message: AxisArray) -> None:
                 ch_data[idx]["bank"] = chr(ord("A") + bank - 1)
                 ch_data[idx]["elec"] = term
                 ch_data[idx]["headstage"] = entry.headstage
+                ch_data[idx]["array"] = _array_identity(entry.label, entry.headstage, chr(ord("A") + bank - 1))
                 cmp_mask[idx] = True
 
         self.state.channel_axis = CoordinateAxis(data=ch_data, dims=["ch"], unit="struct")
@@ -190,7 +236,7 @@ def _apply_incoming_positions(message: AxisArray, ch_data: np.ndarray, n_total:
         if incoming is None or not names or not ({"x", "y"} <= set(names)):
             return src_mask
 
-        copy_fields = [f for f in ("x", "y", "size", "bank", "elec", "headstage") if f in names]
+        copy_fields = [f for f in ("x", "y", "size", "bank", "elec", "headstage", "array") if f in names]
         seen_origin = False
         for i in range(min(n_total, incoming.shape[0])):
             for f in copy_fields:
@@ -237,9 +283,13 @@ def _fill_auto_grid(self) -> None:
             ch_data[idx]["x"] = (i % grid_size) * step
             ch_data[idx]["y"] = start_row + (i // grid_size) * step
             ch_data[idx]["size"] = step  # synthetic electrodes sized to the grid pitch
-            ch_data[idx]["bank"] = chr(next_bank_ord + i // 32)
+            auto_bank = chr(next_bank_ord + i // 32)
+            ch_data[idx]["bank"] = auto_bank
             ch_data[idx]["elec"] = (i % 32) + 1
             ch_data[idx]["headstage"] = 0  # auto-grid channels have no headstage
+            # No label to derive an array from, so this falls back to the
+            # synthetic bank -- 32-channel clusters, as if grouped by bank.
+            ch_data[idx]["array"] = _array_identity("", 0, auto_bank)
 
     @staticmethod
     def _placed_pitch(ch_data: np.ndarray, placed_mask: np.ndarray) -> int:

tests/test_channel_map.py

@@ -11,6 +11,7 @@
     ChannelMapProcessor,
     ChannelMapSettings,
     ChannelMapUnitSettings,
+    _array_identity,
 )
 
 CMP_FILE = str(pathlib.Path(__file__).resolve().parent / "128ChannelDefaultMapping.cmp")
@@ -438,3 +439,97 @@ def test_unstructured_incoming_still_auto_grids_from_origin(self):
         assert data[0]["x"] == 0
         assert data[0]["y"] == 0
         assert not proc.state.src_mask.any()
+
+
+# ---------------------------------------------------------------------------
+# array field
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "label, headstage, bank, expected",
+    [
+        ("elec1-m1-63", 1, "A", "hs1-elec1"),
+        ("elec2-dlpfc-128", 1, "C", "hs1-elec2"),
+        # Same label on another headstage must stay a distinct array.
+        ("elec1-m1-63", 2, "A", "hs2-elec1"),
+        # Unknown headstage: identity without the prefix.
+        ("elec1-m1-63", 0, "A", "elec1"),
+        # A pre-region CMP still yields a usable array identity.
+        ("elec1-63", 1, "A", "hs1-elec1"),
+        # No array structure in the label -> fall back to the bank.
+        ("chan5", 1, "B", "hs1-bankB"),
+        ("", 0, "D", "bankD"),
+        # Neither label nor bank.
+        ("chan5", 1, "", ""),
+    ],
+)
+def test_array_identity(label: str, headstage: int, bank: str, expected: str) -> None:
+    assert _array_identity(label, headstage, bank) == expected
+
+
+def test_cmp_populates_array_field() -> None:
+    """A CMP overlay derives ``array`` from the label + headstage."""
+    proc = _make_processor(CMP_FILE, start_chan=1, hs_id=1)
+    out = proc(_make_message(128))
+    ch = out.axes["ch"].data
+
+    assert "array" in ch.dtype.names
+    for i in range(128):
+        assert ch[i]["array"] == _array_identity(ch[i]["label"], ch[i]["headstage"], ch[i]["bank"])
+    # Every mapped channel belongs to some array.
+    assert all(str(v) for v in ch["array"])
+
+
+def test_array_groups_are_larger_than_banks() -> None:
+    """``array`` groups whole electrode arrays; ``bank`` splits them by connector.
+
+    This is the property that makes ``array`` the useful field for per-array
+    rereferencing: a 32-channel bank is a wiring artifact, while the array is
+    the physical implant.
+    """
+    proc = _make_processor(CMP_FILE, start_chan=1, hs_id=1)
+    ch = proc(_make_message(128)).axes["ch"].data
+
+    def sizes(field: str) -> set[int]:
+        counts: dict[str, int] = {}
+        for value in ch[field]:
+            counts[str(value)] = counts.get(str(value), 0) + 1
+        return set(counts.values())
+
+    # This fixture's labels (chan<N>) carry no array structure, so array
+    # falls back to the bank and the groups match exactly.
+    assert max(sizes("array")) >= max(sizes("bank"))
+
+
+def test_two_headstages_do_not_merge_arrays() -> None:
+    """The same CMP on two headstages yields disjoint arrays, not one merged set."""
+    proc = ChannelMapProcessor(
+        settings=ChannelMapUnitSettings(
+            cmp_configs=(
+                ChannelMapSettings(filepath=CMP_FILE, start_chan=1, hs_id=1),
+                ChannelMapSettings(filepath=CMP_FILE, start_chan=129, hs_id=2),
+            )
+        )
+    )
+    ch = proc(_make_message(256)).axes["ch"].data
+
+    hs1 = {str(v) for v in ch["array"][:128]}
+    hs2 = {str(v) for v in ch["array"][128:]}
+    assert hs1 and hs2
+    assert not (hs1 & hs2), f"identical CMPs on two headstages merged into shared arrays: {hs1 & hs2}"
+    assert all(v.startswith("hs1-") for v in hs1)
+    assert all(v.startswith("hs2-") for v in hs2)
+
+
+def test_auto_grid_array_mirrors_bank() -> None:
+    """Unmapped channels get an ``array`` derived from their synthetic bank.
+
+    Without the fallback they would share one empty value and collapse into
+    a single cluster spanning every unmapped channel.
+    """
+    proc = _make_processor()  # no CMP -> pure auto-grid
+    ch = proc(_make_message(64)).axes["ch"].data
+
+    for i in range(64):
+        assert ch[i]["array"] == f"bank{ch[i]['bank']}"