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 e6cbf13..a720239 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/src/ezmsg/learn/process/ssr.py @@ -58,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 @@ -71,6 +72,7 @@ # callers need to tune it. MIN_REREF_CLUSTER_SIZE = 3 + # --------------------------------------------------------------------------- # Base: Self-supervised regression # --------------------------------------------------------------------------- @@ -206,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( @@ -214,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 ------------------------------------------------------ @@ -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) @@ -374,6 +378,14 @@ class LRRSettings(SelfSupervisedRegressionSettings): """Passed to :class:`AffineTransformTransformer` for the block-diagonal merge threshold.""" + 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 (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 class LRRState(SelfSupervisedRegressionState): @@ -430,9 +442,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) - 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 79904e1..ac99ff8 100644 --- a/tests/unit/test_ssr.py +++ b/tests/unit/test_ssr.py @@ -10,6 +10,7 @@ MIN_REREF_CLUSTER_SIZE, LRRSettings, LRRTransformer, + RereferenceKind, ) # --------------------------------------------------------------------------- @@ -516,3 +517,172 @@ 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=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) + + 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=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) + + 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=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) + + 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=RereferenceKind.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=RereferenceKind.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) + + +# --------------------------------------------------------------------------- +# 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)