LRR: pass clusters below MIN_REREF_CLUSTER_SIZE through untouched

src/ezmsg/learn/process/ssr.py

@@ -61,6 +61,16 @@
 from ezmsg.sigproc.util.channels import channel_clusters_from_field
 from ezmsg.util.messages.axisarray import AxisArray
 
+# Minimum channels a cluster needs before it is rereferenced. Rereferencing
+# regresses each channel against the *others* in its cluster, so a cluster with
+# fewer than this many channels has too few references to be meaningful (1 -> no
+# reference at all; 2 -> a single, degenerate reference). Such clusters are passed
+# through untouched (identity). This also makes sliced/partial inputs robust: a
+# cluster reduced to a channel or two (or an empty cluster) is a no-op rather than
+# a crash or an unstable fit. Kept a module const for now; promote to a setting if
+# callers need to tune it.
+MIN_REREF_CLUSTER_SIZE = 3
+
 # ---------------------------------------------------------------------------
 # Base: Self-supervised regression
 # ---------------------------------------------------------------------------
@@ -191,7 +201,19 @@ def _validate_clusters(self, n_channels: int) -> None:
         """Raise if any cluster index is out of range."""
         clusters = self._get_channel_clusters(n_channels)
         if clusters is None:
-            return
+            return  # implicit single cluster
+        if len(clusters) == 0:
+            # 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.)
+            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})")
@@ -227,7 +249,10 @@ def _solve_weights(self, cxx):
 
         for cluster in clusters:
             k = len(cluster)
-            if k <= 1:
+            if k < MIN_REREF_CLUSTER_SIZE:
+                # Too few channels to rereference against -- leave these channels
+                # untouched (W rows stay 0 -> identity). Covers sliced/partial
+                # clusters down to a single channel; never raises.
                 continue
 
             idx_xp = xp.asarray(cluster) if dev is None else xp.asarray(cluster, device=dev)
@@ -289,6 +314,10 @@ def partial_fit(self, message: AxisArray) -> None:  # type: ignore[override]
             data = xp.permute_dims(data, perm)
 
         n_channels = data.shape[-1]
+        if n_channels == 0:
+            # No channels to fit (e.g. a fully sliced-out hub). Leave the weights
+            # untouched; _process passes the 0-channel data through unchanged.
+            return
         X = xp.reshape(data, (-1, n_channels))
 
         # Covariance stays in the source namespace for accumulation.
@@ -309,6 +338,9 @@ def fit(self, X: np.ndarray) -> None:
         """Batch fit from a raw numpy array (samples x channels)."""
         n_channels = X.shape[-1]
         self._validate_clusters(n_channels)
+        if n_channels == 0:
+            # No channels to fit -- same 0-channel no-op as partial_fit.
+            return
         X = np.asarray(X, dtype=np.float64).reshape(-1, n_channels)
         self._state.cxx = X.T @ X
         self._state.n_samples = X.shape[0]
@@ -388,8 +420,13 @@ def _on_weights_updated(self) -> None:
     # -- transform -----------------------------------------------------------
 
     def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis or message.dims[-1]
+        if message.data.shape[message.get_axis_idx(axis)] == 0:
+            # No channels (e.g. a fully sliced-out hub): nothing to rereference.
+            # Pass the 0-channel message through unchanged -- building an affine
+            # from empty channel clusters would raise downstream.
+            return message
         if self._state.affine is None:
-            axis = self.settings.axis or message.dims[-1]
             axis_idx = message.get_axis_idx(axis)
             n_channels = message.data.shape[axis_idx]
 

tests/unit/test_ssr.py

@@ -6,7 +6,11 @@
 import pytest
 from ezmsg.util.messages.axisarray import AxisArray
 
-from ezmsg.learn.process.ssr import LRRSettings, LRRTransformer
+from ezmsg.learn.process.ssr import (
+    MIN_REREF_CLUSTER_SIZE,
+    LRRSettings,
+    LRRTransformer,
+)
 
 # ---------------------------------------------------------------------------
 # Helpers
@@ -435,3 +439,80 @@ def test_precalculated_weights_from_file(self):
 
         expected = X @ (np.eye(n_ch) - W)
         np.testing.assert_allclose(out.data, expected, atol=1e-10)
+
+
+def _common_mode_data(n_times: int = 400, n_ch: int = 8, rng=None) -> np.ndarray:
+    """Noise plus a shared common-mode component, so rereferencing (which
+    regresses out shared signal) produces a clearly non-identity output."""
+    if rng is None:
+        rng = np.random.default_rng(7)
+    common = rng.standard_normal((n_times, 1))
+    return rng.standard_normal((n_times, n_ch)) + common
+
+
+class TestLowChannelPassthrough:
+    """Clusters smaller than MIN_REREF_CLUSTER_SIZE (and empty inputs) pass
+    through untouched instead of crashing -- so sliced/partial channel sets are
+    safe (e.g. a hub left with no channels after an upstream region slice)."""
+
+    def _fit_process(self, data: np.ndarray, banks: list[str]) -> np.ndarray:
+        proc = LRRTransformer(LRRSettings(axis="ch", cluster_by_field="bank"))
+        for _ in range(8):
+            proc.partial_fit(_banked_axisarray(data, banks))
+        return np.asarray(proc(_banked_axisarray(data, banks)).data)
+
+    def test_zero_channels_passthrough(self):
+        """0 channels (fully sliced-out hub) must not crash on fit or process."""
+        proc = LRRTransformer(LRRSettings(axis="ch", cluster_by_field="bank"))
+        empty = _banked_axisarray(np.zeros((10, 0)), [])
+        proc.partial_fit(empty)  # no channels to fit -- must be a no-op
+        out = proc(empty)  # must pass through, not build an affine from []
+        assert out.data.shape == (10, 0)
+
+    def test_zero_channels_batch_fit(self):
+        """Batch fit() with 0 channels is the same no-op as partial_fit."""
+        proc = LRRTransformer(LRRSettings(axis="ch"))
+        proc.fit(np.zeros((10, 0)))
+        out = proc(_banked_axisarray(np.zeros((10, 0)), []))
+        assert out.data.shape == (10, 0)
+
+    def test_single_channel_identity(self):
+        rng = np.random.default_rng(1)
+        X = _common_mode_data(n_ch=1, rng=rng)
+        out = self._fit_process(X, ["A"])
+        np.testing.assert_allclose(out, X, atol=1e-10)
+
+    def test_below_threshold_identity(self):
+        """A cluster with < MIN_REREF_CLUSTER_SIZE channels is left untouched."""
+        n = MIN_REREF_CLUSTER_SIZE - 1
+        rng = np.random.default_rng(2)
+        X = _common_mode_data(n_ch=n, rng=rng)
+        out = self._fit_process(X, ["A"] * n)
+        np.testing.assert_allclose(out, X, atol=1e-10)
+
+    def test_at_threshold_rereferences(self):
+        """A cluster with exactly MIN_REREF_CLUSTER_SIZE channels is rereferenced."""
+        n = MIN_REREF_CLUSTER_SIZE
+        rng = np.random.default_rng(3)
+        X = _common_mode_data(n_ch=n, rng=rng)
+        out = self._fit_process(X, ["A"] * n)
+        assert np.max(np.abs(out - X)) > 1e-3
+
+    def test_mixed_small_and_large_clusters(self):
+        """Per-cluster: a full bank rereferences while a lone-channel bank in the
+        same message passes through untouched."""
+        big = MIN_REREF_CLUSTER_SIZE + 1
+        rng = np.random.default_rng(4)
+        X = _common_mode_data(n_ch=big + 1, rng=rng)
+        banks = ["A"] * big + ["B"]  # bank A: big ch, bank B: 1 ch
+        out = self._fit_process(X, banks)
+        np.testing.assert_allclose(out[:, big], X[:, big], atol=1e-10)  # lone B ch untouched
+        assert np.max(np.abs(out[:, :big] - X[:, :big])) > 1e-3  # bank A rereferenced
+
+    def test_empty_explicit_clusters_with_channels_raises(self):
+        """channel_clusters=[] with real channels is a misconfiguration: fail fast
+        rather than silently disable rereferencing (the empty list is only
+        tolerated when there are no channels)."""
+        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)))