Add LRRSettings.init_default (identity|CAR) for cold-start rereferencing

src/ezmsg/learn/process/ssr.py

@@ -39,6 +39,7 @@
 
 from __future__ import annotations
 
+import enum
 import os
 import typing
 from abc import abstractmethod
@@ -71,6 +72,23 @@
 # callers need to tune it.
 MIN_REREF_CLUSTER_SIZE = 3
 
+
+class RereferenceInit(str, enum.Enum):
+    """Effective transform to apply before any weights are set.
+
+    Governs the cold-start behavior of an affine rereference (``weights=None`` and
+    no fit yet). ``str`` enum so it round-trips through config as its plain value.
+    """
+
+    IDENTITY = "identity"
+    """Pass the signal through unchanged (legacy default)."""
+
+    CAR = "car"
+    """Per-cluster leave-one-out common-average reference, derived from the
+    resolved channel clusters. A useful cold start (e.g. before any fitted LRR
+    weights exist); replaced as soon as weights are provided or fit."""
+
+
 # ---------------------------------------------------------------------------
 # Base: Self-supervised regression
 # ---------------------------------------------------------------------------
@@ -374,6 +392,12 @@ class LRRSettings(SelfSupervisedRegressionSettings):
     """Passed to :class:`AffineTransformTransformer` for the block-diagonal
     merge threshold."""
 
+    init_default: RereferenceInit = RereferenceInit.IDENTITY
+    """Effective transform used when ``weights`` is None and nothing has been fit
+    yet. ``IDENTITY`` passes through (legacy); ``CAR`` applies per-cluster
+    leave-one-out common-average referencing from the resolved clusters. Provided
+    or fitted weights always take precedence over this cold-start default."""
+
 
 @processor_state
 class LRRState(SelfSupervisedRegressionState):
@@ -417,6 +441,39 @@ def _on_weights_updated(self) -> None:
                 )
             )
 
+    # -- cold-start default (no weights yet) --------------------------------
+
+    def _car_effective_matrix(self, n_channels: int, xp, dtype, dev):
+        """Per-cluster leave-one-out CAR as an effective ``I - W`` matrix.
+
+        Within a cluster of ``k >= MIN_REREF_CLUSTER_SIZE`` channels the sub-block
+        is ``y_i = x_i - mean_{j != i} x_j`` (diagonal 1, off-diagonal
+        ``-1/(k-1)``); clusters smaller than that -- and all cross-cluster terms --
+        stay identity, matching the fit's passthrough for tiny/sliced clusters.
+        Built in the message's array namespace via a selection-matrix scatter, so
+        GPU-backed arrays stay on device (mirrors :meth:`_solve_weights`).
+        """
+        eye_n = xp_create(xp.eye, n_channels, dtype=dtype, device=dev)
+        effective = eye_n
+        clusters = self._get_channel_clusters(n_channels)
+        if clusters is None:
+            clusters = [list(range(n_channels))]
+        for cluster in clusters:
+            k = len(cluster)
+            if k < MIN_REREF_CLUSTER_SIZE:
+                continue  # too few references: leave this block identity
+            idx = xp.asarray(cluster) if dev is None else xp.asarray(cluster, device=dev)
+            eye_k = xp_create(xp.eye, k, dtype=dtype, device=dev)
+            ones_k = xp_create(xp.ones, (k, k), dtype=dtype, device=dev)
+            # leave-one-out CAR sub-block: (k/(k-1)) I_k - (1/(k-1)) J_k
+            block = (k / (k - 1.0)) * eye_k - (1.0 / (k - 1.0)) * ones_k
+            # scatter (block - I_k) onto the identity so the sub-block becomes block
+            S = xp.take(eye_n, idx, axis=1)  # (n, k)
+            effective = effective + xp.matmul(
+                S, xp.matmul(block - eye_k, xp.permute_dims(S, (1, 0)))
+            )
+        return effective
+
     # -- transform -----------------------------------------------------------
 
     def _process(self, message: AxisArray) -> AxisArray:
@@ -432,7 +489,15 @@ def _process(self, message: AxisArray) -> AxisArray:
 
             xp = get_namespace(message.data)
             dev = array_device(message.data)
-            effective = xp_create(xp.eye, n_channels, dtype=message.data.dtype, device=dev)
+            # No weights provided or fit yet: use the configured cold-start default.
+            if self.settings.init_default == RereferenceInit.CAR:
+                effective = self._car_effective_matrix(
+                    n_channels, xp, message.data.dtype, dev
+                )
+            else:
+                effective = xp_create(
+                    xp.eye, n_channels, dtype=message.data.dtype, device=dev
+                )
             self._state.affine = AffineTransformTransformer(
                 AffineTransformSettings(
                     weights=effective,

tests/unit/test_ssr.py

@@ -10,6 +10,7 @@
     MIN_REREF_CLUSTER_SIZE,
     LRRSettings,
     LRRTransformer,
+    RereferenceInit,
 )
 
 # ---------------------------------------------------------------------------
@@ -516,3 +517,97 @@ def test_empty_explicit_clusters_with_channels_raises(self):
         proc = LRRTransformer(LRRSettings(axis="ch", channel_clusters=[]))
         with pytest.raises(ValueError, match="empty but the input has"):
             proc.partial_fit(_make_axisarray(_random_data(n_ch=8)))
+
+
+class TestCARInit:
+    """init_default=CAR: cold-start per-cluster leave-one-out CAR when there are
+    no weights and nothing has been fit."""
+
+    @staticmethod
+    def _loo_car(X: np.ndarray, clusters) -> np.ndarray:
+        """Reference per-cluster leave-one-out CAR: y_i = x_i - mean_{j!=i} x_j."""
+        out = X.copy()
+        for cl in clusters:
+            if len(cl) < MIN_REREF_CLUSTER_SIZE:
+                continue
+            block = X[:, cl]
+            loo = (block.sum(axis=1, keepdims=True) - block) / (len(cl) - 1)
+            out[:, cl] = block - loo
+        return out
+
+    def test_car_applies_leave_one_out_per_cluster(self):
+        clusters = [[0, 1, 2, 3], [4, 5, 6, 7]]
+        X = _random_data(n_ch=8)
+        proc = LRRTransformer(
+            LRRSettings(channel_clusters=clusters, init_default=RereferenceInit.CAR)
+        )
+        out = proc.send(_make_axisarray(X))  # no fit / no weights
+        np.testing.assert_allclose(out.data, self._loo_car(X, clusters), atol=1e-10)
+
+    def test_car_leaves_small_clusters_identity(self):
+        # first cluster (size 2 < MIN_REREF_CLUSTER_SIZE) must pass through
+        clusters = [[0, 1], [2, 3, 4, 5, 6, 7]]
+        X = _random_data(n_ch=8)
+        proc = LRRTransformer(
+            LRRSettings(channel_clusters=clusters, init_default=RereferenceInit.CAR)
+        )
+        out = proc.send(_make_axisarray(X))
+        np.testing.assert_allclose(out.data[:, :2], X[:, :2], atol=1e-12)
+        np.testing.assert_allclose(out.data, self._loo_car(X, clusters), atol=1e-10)
+
+    def test_car_from_bank_field(self):
+        """cluster_by_field='bank' + CAR reproduces per-bank leave-one-out CAR."""
+        n_ch = 8
+        ch = np.zeros(n_ch, dtype=[("bank", "U1")])
+        ch["bank"][:4], ch["bank"][4:] = "A", "B"
+        X = _random_data(n_ch=n_ch)
+        msg = AxisArray(
+            data=X,
+            dims=["time", "ch"],
+            axes={
+                "time": AxisArray.TimeAxis(fs=100.0, offset=0.0),
+                "ch": AxisArray.CoordinateAxis(data=ch, dims=["ch"]),
+            },
+            key="test",
+        )
+        proc = LRRTransformer(
+            LRRSettings(axis="ch", cluster_by_field="bank", init_default=RereferenceInit.CAR)
+        )
+        out = proc.send(msg)
+        np.testing.assert_allclose(
+            out.data, self._loo_car(X, [[0, 1, 2, 3], [4, 5, 6, 7]]), atol=1e-10
+        )
+
+    def test_default_init_is_identity_passthrough(self):
+        """Default (IDENTITY) with no weights is unchanged legacy passthrough."""
+        X = _random_data(n_ch=8)
+        proc = LRRTransformer(LRRSettings(channel_clusters=[[0, 1, 2, 3], [4, 5, 6, 7]]))
+        out = proc.send(_make_axisarray(X))
+        np.testing.assert_allclose(out.data, X, atol=1e-12)
+
+    def test_provided_weights_override_car(self):
+        """Explicit weights win over the CAR cold-start default."""
+        X = _random_data(n_ch=8)
+        # W = 0 => effective I - W = identity, so output is passthrough (not CAR).
+        proc = LRRTransformer(
+            LRRSettings(weights=np.zeros((8, 8)), init_default=RereferenceInit.CAR)
+        )
+        out = proc.send(_make_axisarray(X))
+        np.testing.assert_allclose(out.data, X, atol=1e-12)
+
+    def test_fit_overrides_car(self):
+        """A fitted LRR takes precedence over the CAR cold-start default: once
+        weights are learned, output is the fitted rereference, not CAR."""
+        clusters = [[0, 1, 2, 3], [4, 5, 6, 7]]
+        X = _random_data(n_ch=8, n_times=400)
+        msg = _make_axisarray(X)
+        proc = LRRTransformer(
+            LRRSettings(channel_clusters=clusters, init_default=RereferenceInit.CAR)
+        )
+        proc.partial_fit(msg)
+        out = proc.send(msg)
+
+        fitted = X @ (np.eye(8) - proc.state.weights)
+        np.testing.assert_allclose(out.data, fitted, atol=1e-8)
+        # And it is NOT the CAR cold-start.
+        assert not np.allclose(out.data, self._loo_car(X, clusters), atol=1e-8)