From a08782fbd2314040b6a7d5ced17e047d9052ac1c Mon Sep 17 00:00:00 2001 From: kylmcgr Date: Thu, 23 Jul 2026 12:54:39 -0600 Subject: [PATCH 1/3] Add LRRSettings.init_default (identity|CAR) for cold-start rereferencing --- src/ezmsg/learn/process/ssr.py | 67 +++++++++++++++++++++++- tests/unit/test_ssr.py | 95 ++++++++++++++++++++++++++++++++++ 2 files changed, 161 insertions(+), 1 deletion(-) diff --git a/src/ezmsg/learn/process/ssr.py b/src/ezmsg/learn/process/ssr.py index e6cbf13..33ae1d5 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/src/ezmsg/learn/process/ssr.py @@ -39,6 +39,7 @@ from __future__ import annotations +import enum import os import typing from abc import abstractmethod @@ -71,6 +72,23 @@ # callers need to tune it. MIN_REREF_CLUSTER_SIZE = 3 + +class RereferenceInit(str, enum.Enum): + """Effective transform to apply before any weights are set. + + Governs the cold-start behavior of an affine rereference (``weights=None`` and + no fit yet). ``str`` enum so it round-trips through config as its plain value. + """ + + IDENTITY = "identity" + """Pass the signal through unchanged (legacy default).""" + + CAR = "car" + """Per-cluster leave-one-out common-average reference, derived from the + resolved channel clusters. A useful cold start (e.g. before any fitted LRR + weights exist); replaced as soon as weights are provided or fit.""" + + # --------------------------------------------------------------------------- # Base: Self-supervised regression # --------------------------------------------------------------------------- @@ -374,6 +392,12 @@ class LRRSettings(SelfSupervisedRegressionSettings): """Passed to :class:`AffineTransformTransformer` for the block-diagonal merge threshold.""" + init_default: RereferenceInit = RereferenceInit.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. Provided + or fitted weights always take precedence over this cold-start default.""" + @processor_state class LRRState(SelfSupervisedRegressionState): @@ -417,6 +441,39 @@ def _on_weights_updated(self) -> None: ) ) + # -- cold-start default (no weights yet) -------------------------------- + + def _car_effective_matrix(self, n_channels: int, xp, dtype, dev): + """Per-cluster leave-one-out CAR as an effective ``I - W`` matrix. + + Within a cluster of ``k >= MIN_REREF_CLUSTER_SIZE`` channels the sub-block + is ``y_i = x_i - mean_{j != i} x_j`` (diagonal 1, off-diagonal + ``-1/(k-1)``); clusters smaller than that -- and all cross-cluster terms -- + stay identity, matching the fit's passthrough for tiny/sliced clusters. + Built in the message's array namespace via a selection-matrix scatter, so + GPU-backed arrays stay on device (mirrors :meth:`_solve_weights`). + """ + eye_n = xp_create(xp.eye, n_channels, dtype=dtype, device=dev) + effective = eye_n + clusters = self._get_channel_clusters(n_channels) + if clusters is None: + clusters = [list(range(n_channels))] + for cluster in clusters: + k = len(cluster) + if k < MIN_REREF_CLUSTER_SIZE: + continue # too few references: leave this block identity + idx = xp.asarray(cluster) if dev is None else xp.asarray(cluster, device=dev) + eye_k = xp_create(xp.eye, k, dtype=dtype, device=dev) + ones_k = xp_create(xp.ones, (k, k), dtype=dtype, device=dev) + # leave-one-out CAR sub-block: (k/(k-1)) I_k - (1/(k-1)) J_k + block = (k / (k - 1.0)) * eye_k - (1.0 / (k - 1.0)) * ones_k + # scatter (block - I_k) onto the identity so the sub-block becomes block + S = xp.take(eye_n, idx, axis=1) # (n, k) + effective = effective + xp.matmul( + S, xp.matmul(block - eye_k, xp.permute_dims(S, (1, 0))) + ) + return effective + # -- transform ----------------------------------------------------------- def _process(self, message: AxisArray) -> AxisArray: @@ -432,7 +489,15 @@ def _process(self, message: AxisArray) -> AxisArray: 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: use the configured cold-start default. + if self.settings.init_default == RereferenceInit.CAR: + effective = self._car_effective_matrix( + n_channels, xp, message.data.dtype, dev + ) + else: + effective = xp_create( + xp.eye, n_channels, dtype=message.data.dtype, device=dev + ) self._state.affine = AffineTransformTransformer( AffineTransformSettings( weights=effective, diff --git a/tests/unit/test_ssr.py b/tests/unit/test_ssr.py index 79904e1..c173c04 100644 --- a/tests/unit/test_ssr.py +++ b/tests/unit/test_ssr.py @@ -10,6 +10,7 @@ MIN_REREF_CLUSTER_SIZE, LRRSettings, LRRTransformer, + RereferenceInit, ) # --------------------------------------------------------------------------- @@ -516,3 +517,97 @@ 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=RereferenceInit.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=RereferenceInit.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=RereferenceInit.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=RereferenceInit.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=RereferenceInit.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) From a8d4279ba260e92555caa52ea6021b9ede993cee Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Thu, 23 Jul 2026 19:11:48 -0400 Subject: [PATCH 2/3] Delegate CAR cold-start matrix to ezmsg-sigproc util.rereference ezmsg-sigproc 2.34.0 (ezmsg-org/ezmsg-sigproc#186) hoists deterministic cluster-aware rereference matrix construction into ezmsg.sigproc.util.rereference. Use it here: - Replace the local RereferenceInit enum with the imported RereferenceKind (re-exported from this module; same "identity"/"car" config values). - Delete _car_effective_matrix: the cold-start branch in _process is now a single rereference_matrix(...) call. The matrix is built as host-side numpy; AffineTransformTransformer converts it to the message's namespace/dtype/device on first use, so the on-device selection-matrix scatter is no longer needed. - Use the shared validate_channel_clusters() from util.channels for cluster bounds checking (empty-list policy stays here). - Bump ezmsg-sigproc floor to 2.34.0 (and test dep ezmsg-simbiophys to 1.8.0). Behavior is unchanged for float streams; integer-dtype streams now get float64 output from the cold-start transform (weights stay float64 instead of being cast to the message's integer dtype). Co-Authored-By: Claude Fable 5 --- pyproject.toml | 4 +- src/ezmsg/learn/process/ssr.py | 90 ++++++++-------------------------- tests/unit/test_ssr.py | 26 +++------- 3 files changed, 30 insertions(+), 90 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b8f53ae..5fa145f 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/src/ezmsg/learn/process/ssr.py b/src/ezmsg/learn/process/ssr.py index 33ae1d5..f9497b0 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/src/ezmsg/learn/process/ssr.py @@ -39,7 +39,6 @@ from __future__ import annotations -import enum import os import typing from abc import abstractmethod @@ -59,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 @@ -73,22 +73,6 @@ MIN_REREF_CLUSTER_SIZE = 3 -class RereferenceInit(str, enum.Enum): - """Effective transform to apply before any weights are set. - - Governs the cold-start behavior of an affine rereference (``weights=None`` and - no fit yet). ``str`` enum so it round-trips through config as its plain value. - """ - - IDENTITY = "identity" - """Pass the signal through unchanged (legacy default).""" - - CAR = "car" - """Per-cluster leave-one-out common-average reference, derived from the - resolved channel clusters. A useful cold start (e.g. before any fitted LRR - weights exist); replaced as soon as weights are provided or fit.""" - - # --------------------------------------------------------------------------- # Base: Self-supervised regression # --------------------------------------------------------------------------- @@ -224,7 +208,7 @@ 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( @@ -232,9 +216,7 @@ def _validate_clusters(self, n_channels: int) -> None: "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 ------------------------------------------------------ @@ -392,11 +374,13 @@ class LRRSettings(SelfSupervisedRegressionSettings): """Passed to :class:`AffineTransformTransformer` for the block-diagonal merge threshold.""" - init_default: RereferenceInit = RereferenceInit.IDENTITY + 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. Provided - or fitted weights always take precedence over this cold-start default.""" + 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 @@ -441,39 +425,6 @@ def _on_weights_updated(self) -> None: ) ) - # -- cold-start default (no weights yet) -------------------------------- - - def _car_effective_matrix(self, n_channels: int, xp, dtype, dev): - """Per-cluster leave-one-out CAR as an effective ``I - W`` matrix. - - Within a cluster of ``k >= MIN_REREF_CLUSTER_SIZE`` channels the sub-block - is ``y_i = x_i - mean_{j != i} x_j`` (diagonal 1, off-diagonal - ``-1/(k-1)``); clusters smaller than that -- and all cross-cluster terms -- - stay identity, matching the fit's passthrough for tiny/sliced clusters. - Built in the message's array namespace via a selection-matrix scatter, so - GPU-backed arrays stay on device (mirrors :meth:`_solve_weights`). - """ - eye_n = xp_create(xp.eye, n_channels, dtype=dtype, device=dev) - effective = eye_n - clusters = self._get_channel_clusters(n_channels) - if clusters is None: - clusters = [list(range(n_channels))] - for cluster in clusters: - k = len(cluster) - if k < MIN_REREF_CLUSTER_SIZE: - continue # too few references: leave this block identity - idx = xp.asarray(cluster) if dev is None else xp.asarray(cluster, device=dev) - eye_k = xp_create(xp.eye, k, dtype=dtype, device=dev) - ones_k = xp_create(xp.ones, (k, k), dtype=dtype, device=dev) - # leave-one-out CAR sub-block: (k/(k-1)) I_k - (1/(k-1)) J_k - block = (k / (k - 1.0)) * eye_k - (1.0 / (k - 1.0)) * ones_k - # scatter (block - I_k) onto the identity so the sub-block becomes block - S = xp.take(eye_n, idx, axis=1) # (n, k) - effective = effective + xp.matmul( - S, xp.matmul(block - eye_k, xp.permute_dims(S, (1, 0))) - ) - return effective - # -- transform ----------------------------------------------------------- def _process(self, message: AxisArray) -> AxisArray: @@ -487,17 +438,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) - # No weights provided or fit yet: use the configured cold-start default. - if self.settings.init_default == RereferenceInit.CAR: - effective = self._car_effective_matrix( - n_channels, xp, message.data.dtype, dev - ) - else: - 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, diff --git a/tests/unit/test_ssr.py b/tests/unit/test_ssr.py index c173c04..2bf1c3d 100644 --- a/tests/unit/test_ssr.py +++ b/tests/unit/test_ssr.py @@ -10,7 +10,7 @@ MIN_REREF_CLUSTER_SIZE, LRRSettings, LRRTransformer, - RereferenceInit, + RereferenceKind, ) # --------------------------------------------------------------------------- @@ -538,9 +538,7 @@ def _loo_car(X: np.ndarray, clusters) -> np.ndarray: 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=RereferenceInit.CAR) - ) + 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) @@ -548,9 +546,7 @@ 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=RereferenceInit.CAR) - ) + 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) @@ -570,13 +566,9 @@ def test_car_from_bank_field(self): }, key="test", ) - proc = LRRTransformer( - LRRSettings(axis="ch", cluster_by_field="bank", init_default=RereferenceInit.CAR) - ) + 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 - ) + 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.""" @@ -589,9 +581,7 @@ 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=RereferenceInit.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) @@ -601,9 +591,7 @@ def test_fit_overrides_car(self): 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=RereferenceInit.CAR) - ) + proc = LRRTransformer(LRRSettings(channel_clusters=clusters, init_default=RereferenceKind.CAR)) proc.partial_fit(msg) out = proc.send(msg) From 373d9873df921b76e2746b5dc944981982e5ee04 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Thu, 23 Jul 2026 19:41:19 -0400 Subject: [PATCH 3/3] Add LRR backend-preservation tests; fix MLX fit via CPU-stream linalg New TestBackendPreservation (parametrized over mlx and torch, with importorskip guards for platforms without mlx) asserts that: - the output stays in the input's array namespace, - fitted state (cxx, weights) and the internal affine's weight arrays live in that namespace, and - numpy matrices (cold-start CAR/identity, user-provided settings.weights) are converted to the message's backend on first use, with values matching the numpy reference path in all three cases. The tests exposed a pre-existing bug: MLX's linalg.inv/pinv only run on the CPU stream, so partial_fit crashed for mlx-backed messages. _solve_weights now passes stream=mx.cpu for the mlx namespace -- with unified memory this is a scheduling hint, not a host copy, and results remain mlx arrays. Co-Authored-By: Claude Fable 5 --- src/ezmsg/learn/process/ssr.py | 8 +++- tests/unit/test_ssr.py | 87 ++++++++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+), 2 deletions(-) diff --git a/src/ezmsg/learn/process/ssr.py b/src/ezmsg/learn/process/ssr.py index f9497b0..a720239 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/src/ezmsg/learn/process/ssr.py @@ -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) diff --git a/tests/unit/test_ssr.py b/tests/unit/test_ssr.py index 2bf1c3d..ac99ff8 100644 --- a/tests/unit/test_ssr.py +++ b/tests/unit/test_ssr.py @@ -599,3 +599,90 @@ def test_fit_overrides_car(self): 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)