diff --git a/src/ezmsg/learn/process/ssr.py b/src/ezmsg/learn/process/ssr.py index a720239..2061c41 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/src/ezmsg/learn/process/ssr.py @@ -420,11 +420,17 @@ def _on_weights_updated(self) -> None: if self._state.affine is not None: self._state.affine.set_weights(effective) else: + # channel_clusters=None: let the affine detect blocks from the weight + # matrix itself. self._get_channel_clusters() can fall back to + # block_size here (cluster_by_field isn't resolved until a message + # arrives), and blocks finer than W's real ones make the block-diagonal + # matmul silently overwrite -- rereferencing each block against only + # part of its channels. W is the source of truth for its own blocks. self._state.affine = AffineTransformTransformer( AffineTransformSettings( weights=effective, axis=self.settings.axis, - channel_clusters=self._get_channel_clusters(n), + channel_clusters=None, min_cluster_size=self.settings.min_cluster_size, ) ) diff --git a/tests/unit/test_ssr.py b/tests/unit/test_ssr.py index ac99ff8..152b64f 100644 --- a/tests/unit/test_ssr.py +++ b/tests/unit/test_ssr.py @@ -38,6 +38,15 @@ def _random_data(n_times: int = 200, n_ch: int = 8, rng=None) -> np.ndarray: return rng.standard_normal((n_times, n_ch)) +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 + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -442,13 +451,49 @@ def test_precalculated_weights_from_file(self): 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 TestApplyFollowsWeightBlocks: + """Regression: applying fit/loaded weights must equal ``X @ (I - W)`` no matter + what ``block_size`` / ``cluster_by_field`` the transformer carries. + + The block-diagonal apply optimization must follow the WEIGHT MATRIX's real + block structure, not a cluster hint that may not match it. Previously the + affine was built with ``_get_channel_clusters()``, which falls back to + ``block_size`` when the ``cluster_by_field`` metadata is absent at apply time + (e.g. a processor constructed with weights before any message resolves the + field). When those fallback clusters were FINER than the W's true blocks -- + as when a W is fit over non-contiguous electrode-array groups but applied with + a smaller ``block_size`` -- two input sub-groups of one true block mapped to + the same output indices, and the block-diagonal matmul's assignment silently + OVERWROTE the earlier sub-group. Each true block was then rereferenced against + only a subset of its channels, corrupting the output while looking valid. + """ + + def test_apply_ignores_block_size_finer_than_weight_blocks(self): + rng = np.random.default_rng(123) + n_ch = 128 + # W block-diagonal over two NON-CONTIGUOUS 64-ch groups (like two electrode + # arrays interleaved in channel order: an "array" = two connector banks). + groups = [list(range(0, 32)) + list(range(96, 128)), list(range(32, 96))] + X = _random_data(n_times=600, n_ch=n_ch, rng=rng) + + proc_fit = LRRTransformer(LRRSettings(channel_clusters=groups)) + proc_fit.partial_fit(_make_axisarray(X)) + W = proc_fit.state.weights.copy() + # W really is block-diagonal over the non-contiguous 64-ch groups. + for a in groups: + for b in groups: + if a is b: + continue + np.testing.assert_array_equal(W[np.ix_(a, b)], 0.0) + + # Apply the loaded W with block_size=32 (finer than the W's 64-ch blocks) + # and min_cluster_size=32 so the fallback clusters are NOT merged away. + # These 4x32 clusters do not match the W -> the block-diagonal apply used + # to silently overwrite. The result must still be the faithful X @ (I - W). + proc = LRRTransformer(LRRSettings(weights=W, block_size=32, min_cluster_size=32)) + out = proc.send(_make_axisarray(X)) + expected = X @ (np.eye(n_ch) - W) + np.testing.assert_allclose(out.data, expected, atol=1e-10) class TestLowChannelPassthrough: