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
 # ---------------------------------------------------------------------------
@@ -190,7 +200,10 @@ def _get_channel_clusters(self, n_channels: int) -> list[list[int]] | None:
     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:
+        if not clusters:
+            # None (single implicit cluster) or empty (0 channels / no clustered
+            # data). Nothing to validate -- and np.concatenate([]) would raise,
+            # which is what broke on a fully sliced-out (0-channel) input.
             return
         all_indices = np.concatenate([np.asarray(g) for g in clusters])
         if np.any((all_indices < 0) | (all_indices >= n_channels)):
@@ -227,7 +240,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 +305,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.
@@ -388,8 +408,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,65 @@ 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_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