lrr cluster by field

pyproject.toml

@@ -11,7 +11,7 @@ dynamic = ["version"]
 dependencies = [
     "ezmsg>=3.9.0",
     "ezmsg-baseproc>=1.7.0",
-    "ezmsg-sigproc>=2.23.0",
+    "ezmsg-sigproc>=2.28.0",
     "pandas>=2.2",
     "river>=0.22.0",
     "scikit-learn>=1.6.0",

src/ezmsg/learn/process/ssr.py

@@ -58,6 +58,7 @@
     AffineTransformTransformer,
 )
 from ezmsg.sigproc.util.array import array_device, xp_create
+from ezmsg.sigproc.util.channels import channel_clusters_from_field
 from ezmsg.util.messages.axisarray import AxisArray
 
 # ---------------------------------------------------------------------------
@@ -78,6 +79,13 @@ class SelfSupervisedRegressionSettings(ez.Settings):
     channel_clusters: list[list[int]] | None = None
     """Per-cluster regression.  ``None`` treats all channels as one cluster."""
 
+    cluster_by_field: str | None = None
+    """Derive ``channel_clusters`` automatically from a structured field of the
+    channel coordinate axis (e.g. ``"bank"`` to regress within each electrode
+    bank). Used only when ``channel_clusters`` is None and the axis actually
+    carries that field; otherwise falls back to ``block_size`` (then a single
+    cluster). Explicit ``channel_clusters`` always takes precedence."""
+
     block_size: int | None = None
     """If ``channel_clusters`` is ``None``, use this block size for an automatic clustering."""
 
@@ -94,6 +102,9 @@ class SelfSupervisedRegressionState:
     cxx: object | None = None  # Array API; namespace matches source data.
     n_samples: int = 0
     weights: object | None = None  # Array API; namespace matches cxx.
+    resolved_clusters: list | None = None
+    """Clusters derived from ``cluster_by_field`` at reset (message-dependent),
+    cached so ``_get_channel_clusters`` can return them without the message."""
 
 
 class SelfSupervisedRegressionTransformer(
@@ -115,13 +126,36 @@ class SelfSupervisedRegressionTransformer(
     def _hash_message(self, message: AxisArray) -> int:
         axis = self.settings.axis or message.dims[-1]
         axis_idx = message.get_axis_idx(axis)
-        return hash((message.key, message.data.shape[axis_idx]))
+        components: tuple = (message.key, message.data.shape[axis_idx])
+        # On the cluster_by_field path, re-derive clusters only when the channel
+        # axis gains or loses the target structured field -- a single O(1)
+        # presence boolean rather than O(channels) of field bytes hashed on every
+        # message (which grows with channel count on this hot path). Concession:
+        # if the axis is already structured and the channel count is unchanged, a
+        # change in the field's *values* (a live bank remap) is not detected. That
+        # is safe for real acquisition streams, whose channel->field map is static
+        # for the stream's life; a genuine remap arrives with a new key or a
+        # different channel count, both already folded in above. Mirrors the
+        # ezmsg-sigproc CommonRereference hash.
+        if self.settings.channel_clusters is None and self.settings.cluster_by_field is not None:
+            ax = message.axes.get(axis)
+            names = getattr(getattr(getattr(ax, "data", None), "dtype", None), "names", None)
+            components += (bool(names and self.settings.cluster_by_field in names),)
+        return hash(components)
 
     def _reset_state(self, message: AxisArray) -> None:
         axis = self.settings.axis or message.dims[-1]
         axis_idx = message.get_axis_idx(axis)
         n_channels = message.data.shape[axis_idx]
 
+        # Derive clusters from a structured channel-axis field (e.g. "bank") when
+        # requested and no explicit clusters were given. Cached so the
+        # message-less _get_channel_clusters can return them later.
+        if self.settings.channel_clusters is None and self.settings.cluster_by_field is not None:
+            self._state.resolved_clusters = channel_clusters_from_field(message, axis, self.settings.cluster_by_field)
+        else:
+            self._state.resolved_clusters = None
+
         self._validate_clusters(n_channels)
         self._state.cxx = None
         self._state.n_samples = 0
@@ -141,7 +175,11 @@ def _reset_state(self, message: AxisArray) -> None:
     # -- cluster validation --------------------------------------------------
 
     def _get_channel_clusters(self, n_channels: int) -> list[list[int]] | None:
+        # Precedence: explicit channel_clusters > cluster_by_field-derived
+        # (cached at reset) > block_size > None (single cluster).
         clusters = self.settings.channel_clusters
+        if clusters is None:
+            clusters = getattr(self._state, "resolved_clusters", None)
         if clusters is None and self.settings.block_size is not None:
             clusters = [
                 list(range(i, min(i + self.settings.block_size, n_channels)))

tests/unit/test_ssr.py

@@ -126,6 +126,93 @@ def test_channel_clusters(self):
             assert np.any(off_diag != 0), "Expected non-zero within-cluster weights"
 
 
+def _banked_axisarray(data: np.ndarray, banks: list[str], key: str = "test") -> AxisArray:
+    """AxisArray whose ch axis is a structured CoordinateAxis with a bank field,
+    like ezmsg-blackrock ChannelMap emits."""
+    dt = np.dtype([("label", "U16"), ("bank", "U1"), ("elec", "i4")])
+    ch = np.zeros(len(banks), dtype=dt)
+    ch["bank"] = banks
+    ch["elec"] = list(range(1, len(banks) + 1))
+    ch["label"] = [f"ch{i}" for i in range(len(banks))]
+    return AxisArray(
+        data=data,
+        dims=["time", "ch"],
+        axes={"time": AxisArray.TimeAxis(fs=100.0, offset=0.0), "ch": AxisArray.CoordinateAxis(data=ch, dims=["ch"])},
+        key=key,
+    )
+
+
+class TestClusterByField:
+    def test_bank_field_matches_explicit_clusters(self):
+        """cluster_by_field='bank' derives the same clusters (and weights) as
+        passing the equivalent explicit channel_clusters."""
+        rng = np.random.default_rng(7)
+        banks = ["A", "A", "A", "A", "B", "B", "B", "B"]
+        X = _random_data(n_ch=len(banks), rng=rng)
+
+        proc_field = LRRTransformer(LRRSettings(axis="ch", cluster_by_field="bank"))
+        proc_field.partial_fit(_banked_axisarray(X, banks))
+
+        proc_explicit = LRRTransformer(LRRSettings(axis="ch", channel_clusters=[[0, 1, 2, 3], [4, 5, 6, 7]]))
+        proc_explicit.partial_fit(_make_axisarray(X))
+
+        np.testing.assert_array_equal(proc_field.state.weights, proc_explicit.state.weights)
+        # And cross-bank weights are zero
+        W = proc_field.state.weights
+        np.testing.assert_array_equal(W[np.ix_([0, 1, 2, 3], [4, 5, 6, 7])], 0.0)
+
+    def test_explicit_clusters_take_precedence(self):
+        """Explicit channel_clusters win over cluster_by_field."""
+        rng = np.random.default_rng(8)
+        banks = ["A", "A", "A", "A", "B", "B", "B", "B"]
+        X = _random_data(n_ch=len(banks), rng=rng)
+
+        # One all-channel cluster should override the bank grouping.
+        proc = LRRTransformer(LRRSettings(axis="ch", channel_clusters=[list(range(8))], cluster_by_field="bank"))
+        proc.partial_fit(_banked_axisarray(X, banks))
+        # With a single cluster, cross-"bank" weights are NOT forced to zero.
+        W = proc.state.weights
+        assert np.any(W[np.ix_([0, 1, 2, 3], [4, 5, 6, 7])] != 0)
+
+    def test_missing_field_falls_back_to_block_size(self):
+        """cluster_by_field with no structured bank field falls back to block_size."""
+        rng = np.random.default_rng(9)
+        n_ch = 8
+        X = _random_data(n_ch=n_ch, rng=rng)
+        # Plain axis (no structured bank field) + block_size=4 -> two contiguous blocks.
+        proc_field = LRRTransformer(LRRSettings(axis="ch", cluster_by_field="bank", block_size=4))
+        proc_field.partial_fit(_make_axisarray(X))
+
+        proc_block = LRRTransformer(LRRSettings(axis="ch", block_size=4))
+        proc_block.partial_fit(_make_axisarray(X))
+
+        np.testing.assert_array_equal(proc_field.state.weights, proc_block.state.weights)
+
+    def test_bank_field_value_change_is_not_detected(self):
+        """Intentional concession (mirrors the ezmsg-sigproc CAR fix): a live bank
+        remap at fixed key + channel count is NOT re-derived. ``_hash_message``
+        folds only an O(1) "bank field present" boolean, not the field's bytes, so
+        the per-message hash does not scale with channel count. A genuine remap on
+        real hardware arrives with a new key or channel count (escape hatch below)."""
+        rng = np.random.default_rng(11)
+        X = _random_data(n_ch=4, rng=rng)
+        proc = LRRTransformer(LRRSettings(axis="ch", cluster_by_field="bank"))
+
+        # First arrangement: banks A,A,B,B -> clusters {0,1},{2,3}.
+        proc.partial_fit(_banked_axisarray(X, ["A", "A", "B", "B"], key="x"))
+        assert proc.state.resolved_clusters == [[0, 1], [2, 3]]
+        np.testing.assert_array_equal(proc.state.weights[np.ix_([0, 1], [2, 3])], 0.0)
+
+        # Same key + channel count, different banks -> hash unchanged, so the
+        # cached clusters are (deliberately) NOT re-derived.
+        proc.partial_fit(_banked_axisarray(X, ["A", "B", "A", "B"], key="x"))
+        assert proc.state.resolved_clusters == [[0, 1], [2, 3]]
+
+        # Escape hatch: a new key (as a real remap would carry) forces re-derivation.
+        proc.partial_fit(_banked_axisarray(X, ["A", "B", "A", "B"], key="y"))
+        assert proc.state.resolved_clusters == [[0, 2], [1, 3]]
+
+
 class TestIncrementalAccumulates:
     def test_incremental_accumulates(self):
         """Two partial_fits with incremental=True should match one fit on concatenated data."""