diff --git a/src/ezmsg/learn/process/ssr.py b/src/ezmsg/learn/process/ssr.py index 634f2de..e6cbf13 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/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] diff --git a/tests/unit/test_ssr.py b/tests/unit/test_ssr.py index 6592528..79904e1 100644 --- a/tests/unit/test_ssr.py +++ b/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)))