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

pyproject.toml

@@ -11,7 +11,7 @@ dynamic = ["version"]
 dependencies = [
     "ezmsg>=3.9.0",
     "ezmsg-baseproc>=1.7.0",
-    "ezmsg-sigproc>=2.28.0",
+    "ezmsg-sigproc>=2.34.0",
     "pandas>=2.2",
     "river>=0.22.0",
     "scikit-learn>=1.6.0",
@@ -28,7 +28,7 @@ lint = [
     "ruff>=0.12.9",
 ]
 test = [
-    "ezmsg-simbiophys>=1.3.0",
+    "ezmsg-simbiophys>=1.8.0",
     "hmmlearn>=0.3.3",
     "pytest>=8.4.1",
 ]

src/ezmsg/learn/process/ssr.py

@@ -58,7 +58,8 @@
     AffineTransformTransformer,
 )
 from ezmsg.sigproc.util.array import array_device, xp_create
-from ezmsg.sigproc.util.channels import channel_clusters_from_field
+from ezmsg.sigproc.util.channels import channel_clusters_from_field, validate_channel_clusters
+from ezmsg.sigproc.util.rereference import RereferenceKind, rereference_matrix
 from ezmsg.util.messages.axisarray import AxisArray
 
 # Minimum channels a cluster needs before it is rereferenced. Rereferencing
@@ -71,6 +72,7 @@
 # callers need to tune it.
 MIN_REREF_CLUSTER_SIZE = 3
 
+
 # ---------------------------------------------------------------------------
 # Base: Self-supervised regression
 # ---------------------------------------------------------------------------
@@ -206,17 +208,15 @@ def _validate_clusters(self, n_channels: int) -> None:
             # An empty cluster list is only legitimate with no channels (e.g. a
             # fully sliced-out input). With channels present it means an explicit
             # channel_clusters=[], which would silently disable rereferencing --
-            # fail fast instead. (An empty list also breaks np.concatenate below.)
+            # fail fast instead.
             if n_channels == 0:
                 return
             raise ValueError(
                 f"channel_clusters is empty but the input has {n_channels} channels. "
                 "Pass channel_clusters=None to treat all channels as a single "
                 "cluster, or provide non-empty channel index groups."
             )
-        all_indices = np.concatenate([np.asarray(g) for g in clusters])
-        if np.any((all_indices < 0) | (all_indices >= n_channels)):
-            raise ValueError(f"channel_clusters contains out-of-range indices (valid range: 0..{n_channels - 1})")
+        validate_channel_clusters(clusters, n_channels)
 
     # -- weight solving ------------------------------------------------------
 
@@ -247,6 +247,10 @@ def _solve_weights(self, cxx):
         W = xp_create(xp.zeros, (n, n), dtype=cxx.dtype, device=dev)
         eye_n = xp_create(xp.eye, n, dtype=cxx.dtype, device=dev)
 
+        # MLX linalg ops are CPU-only; with unified memory the explicit CPU
+        # stream is a scheduling hint, not a host copy, and results stay mlx.
+        inv_kwargs = {"stream": xp.cpu} if xp.__name__ == "mlx.core" else {}
+
         for cluster in clusters:
             k = len(cluster)
             if k < MIN_REREF_CLUSTER_SIZE:
@@ -266,9 +270,9 @@ def _solve_weights(self, cxx):
 
             # One inverse per cluster
             try:
-                sub_inv = xp.linalg.inv(sub)
+                sub_inv = xp.linalg.inv(sub, **inv_kwargs)
             except Exception:
-                sub_inv = xp.linalg.pinv(sub)
+                sub_inv = xp.linalg.pinv(sub, **inv_kwargs)
 
             # Diagonal via element-wise product with identity
             diag_vals = xp.sum(sub_inv * eye_k, axis=0)
@@ -374,6 +378,14 @@ class LRRSettings(SelfSupervisedRegressionSettings):
     """Passed to :class:`AffineTransformTransformer` for the block-diagonal
     merge threshold."""
 
+    init_default: RereferenceKind = RereferenceKind.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 (clusters
+    below :data:`MIN_REREF_CLUSTER_SIZE` stay identity, matching the fit's
+    passthrough). Provided or fitted weights always take precedence over this
+    cold-start default."""
+
 
 @processor_state
 class LRRState(SelfSupervisedRegressionState):
@@ -430,9 +442,18 @@ def _process(self, message: AxisArray) -> AxisArray:
             axis_idx = message.get_axis_idx(axis)
             n_channels = message.data.shape[axis_idx]
 
-            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: build the configured cold-start
+            # default (identity, or per-cluster leave-one-out CAR matching the
+            # fit's passthrough for clusters below MIN_REREF_CLUSTER_SIZE).
+            # Built as numpy; the affine transformer converts weights to the
+            # message's namespace/dtype/device on first use.
+            effective = rereference_matrix(
+                self.settings.init_default,
+                n_channels,
+                clusters=self._get_channel_clusters(n_channels),
+                include_current=False,
+                min_reref_size=MIN_REREF_CLUSTER_SIZE,
+            )
             self._state.affine = AffineTransformTransformer(
                 AffineTransformSettings(
                     weights=effective,

tests/unit/test_ssr.py

@@ -10,6 +10,7 @@
     MIN_REREF_CLUSTER_SIZE,
     LRRSettings,
     LRRTransformer,
+    RereferenceKind,
 )
 
 # ---------------------------------------------------------------------------
@@ -516,3 +517,172 @@ 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=RereferenceKind.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=RereferenceKind.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=RereferenceKind.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=RereferenceKind.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=RereferenceKind.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)
+
+
+# ---------------------------------------------------------------------------
+# Backend (array namespace) preservation
+# ---------------------------------------------------------------------------
+
+
+def _backend(name: str):
+    """Return (converter, array_type) for a non-numpy Array API backend,
+    skipping if the library is not installed (e.g. mlx off-macOS)."""
+    if name == "mlx":
+        mx = pytest.importorskip("mlx.core")
+        return mx.array, mx.array
+    torch = pytest.importorskip("torch")
+    return torch.from_numpy, torch.Tensor
+
+
+@pytest.mark.parametrize("backend", ["mlx", "torch"])
+class TestBackendPreservation:
+    """The input's array namespace (mlx / torch) must be preserved to the
+    output, and derived state -- cxx, weights, and the internal affine's
+    weight arrays -- must live in that namespace. Cold-start matrices are
+    deliberately built as numpy and must be converted to the message's
+    backend on first use by the affine transformer."""
+
+    CLUSTERS = [[0, 1, 2, 3], [4, 5, 6, 7]]
+
+    @staticmethod
+    def _affine_weight_arrays(affine):
+        """All weight arrays held by the internal affine (dense or per-cluster)."""
+        if affine.state.weights is not None:
+            return [affine.state.weights]
+        return [sub_w for _, _, sub_w in affine.state.clusters]
+
+    def test_cold_start_car_converts_and_preserves(self, backend):
+        conv, typ = _backend(backend)
+        X = _random_data().astype(np.float32)
+        proc = LRRTransformer(LRRSettings(channel_clusters=self.CLUSTERS, init_default=RereferenceKind.CAR))
+        out = proc.send(_make_axisarray(conv(X.copy())))
+
+        assert isinstance(out.data, typ)
+        weight_arrays = self._affine_weight_arrays(proc.state.affine)
+        assert len(weight_arrays) > 0
+        for w in weight_arrays:
+            assert isinstance(w, typ)
+
+        # Values match the numpy cold-start CAR.
+        ref_proc = LRRTransformer(LRRSettings(channel_clusters=self.CLUSTERS, init_default=RereferenceKind.CAR))
+        ref = ref_proc.send(_make_axisarray(X))
+        np.testing.assert_allclose(np.asarray(out.data), ref.data, atol=1e-5)
+
+    def test_fit_keeps_state_and_output_in_backend(self, backend):
+        conv, typ = _backend(backend)
+        X = _random_data(n_times=400).astype(np.float32)
+        msg = _make_axisarray(conv(X.copy()))
+        proc = LRRTransformer(LRRSettings(channel_clusters=self.CLUSTERS))
+        proc.partial_fit(msg)
+
+        assert isinstance(proc.state.cxx, typ)
+        assert isinstance(proc.state.weights, typ)
+
+        out = proc.send(msg)
+        assert isinstance(out.data, typ)
+        for w in self._affine_weight_arrays(proc.state.affine):
+            assert isinstance(w, typ)
+
+        # Fitted output matches the numpy fit within float32 tolerance.
+        ref_proc = LRRTransformer(LRRSettings(channel_clusters=self.CLUSTERS))
+        ref_proc.partial_fit(_make_axisarray(X))
+        ref = ref_proc.send(_make_axisarray(X))
+        np.testing.assert_allclose(np.asarray(out.data), ref.data, atol=1e-3)
+
+    def test_numpy_settings_weights_with_backend_messages(self, backend):
+        conv, typ = _backend(backend)
+        X = _random_data(n_times=400).astype(np.float32)
+
+        fit_proc = LRRTransformer(LRRSettings(channel_clusters=self.CLUSTERS))
+        fit_proc.partial_fit(_make_axisarray(X))
+        W = np.asarray(fit_proc.state.weights)
+        ref = fit_proc.send(_make_axisarray(X))
+
+        proc = LRRTransformer(LRRSettings(weights=W, channel_clusters=self.CLUSTERS))
+        out = proc.send(_make_axisarray(conv(X.copy())))
+        assert isinstance(out.data, typ)
+        for w in self._affine_weight_arrays(proc.state.affine):
+            assert isinstance(w, typ)
+        np.testing.assert_allclose(np.asarray(out.data), ref.data, atol=1e-3)