diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 6cbe86d..61d1a04 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -34,8 +34,11 @@ jobs: run: uv sync - name: Lint + # uv run, not uv tool run: the latter resolves ruff independently and + # would ignore the version pinned in the lint dependency-group, so a new + # ruff release could turn CI red with no change to this repo. run: - uv tool run ruff check --output-format=github src + uv run ruff check --output-format=github src - name: Run unit tests run: uv run pytest tests/unit tests/dim_reduce -v diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a525e91..8136ed5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.11.12 + rev: v0.16.2 hooks: - id: ruff args: [ --fix ] diff --git a/pyproject.toml b/pyproject.toml index 5fa145f..42b1b3c 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.34.0", + "ezmsg-sigproc>=3.0.0", "pandas>=2.2", "river>=0.22.0", "scikit-learn>=1.6.0", @@ -25,7 +25,7 @@ dev = [ "pre-commit>=4.3.0", ] lint = [ - "ruff>=0.12.9", + "ruff==0.16.2", ] test = [ "ezmsg-simbiophys>=1.8.0", diff --git a/src/ezmsg/learn/collection/sample_adapt_regressor.py b/src/ezmsg/learn/collection/sample_adapt_regressor.py index 1ba5a0f..7611a61 100644 --- a/src/ezmsg/learn/collection/sample_adapt_regressor.py +++ b/src/ezmsg/learn/collection/sample_adapt_regressor.py @@ -82,9 +82,7 @@ class DecodeOutputAdapterProcessor( def _reset_state(self, message: AxisArray) -> None: if self.settings.output_labels is not None: - self.state.ch_axis = AxisArray.CoordinateAxis( - data=np.asarray(self.settings.output_labels), dims=["ch"] - ) + self.state.ch_axis = AxisArray.CoordinateAxis(data=np.asarray(self.settings.output_labels), dims=["ch"]) def _process(self, message: AxisArray) -> AxisArray | None: data = np.asarray(message.data, dtype=float) @@ -97,9 +95,7 @@ def _process(self, message: AxisArray) -> AxisArray | None: ch_axis = self.state.ch_axis else: data = data.reshape((data.shape[0], -1)) if data.ndim > 1 else data.reshape((1, -1)) - ch_axis = AxisArray.CoordinateAxis( - data=np.asarray([f"ch{i}" for i in range(data.shape[-1])]), dims=["ch"] - ) + ch_axis = AxisArray.CoordinateAxis(data=np.asarray([f"ch{i}" for i in range(data.shape[-1])]), dims=["ch"]) # The decoder engines carry a ``time`` axis through (kalman keeps the # input's; the torch path inherits the windower's renamed ``win``->``time`` @@ -266,9 +262,7 @@ def configure(self) -> None: # Window requires zero_pad_until="input" when # window_shift is None (1:1 mode); "none" there only # warns and is coerced to "input". - zero_pad_until="none" - if self.SETTINGS.decode_window_shift is not None - else "input", + zero_pad_until="none" if self.SETTINGS.decode_window_shift is not None else "input", ) ) self.FLATTEN.apply_settings( @@ -293,11 +287,7 @@ def configure(self) -> None: ) ) if needs_adapter: - self.ADAPTER.apply_settings( - DecodeOutputAdapterSettings( - output_labels=self.SETTINGS.output_labels - ) - ) + self.ADAPTER.apply_settings(DecodeOutputAdapterSettings(output_labels=self.SETTINGS.output_labels)) def network(self) -> ez.NetworkDefinition: network = [] diff --git a/src/ezmsg/learn/process/ssr.py b/src/ezmsg/learn/process/ssr.py index 2061c41..18b7c43 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/src/ezmsg/learn/process/ssr.py @@ -6,34 +6,35 @@ :class:`LRRTransformer`. **Framework.** The base class accumulates the channel covariance -``C = X^T X`` and solves per-cluster ridge regressions to obtain a weight +``C = X^T X`` and solves per-group ridge regressions to obtain a weight matrix *W*. Subclasses define what to *do* with *W* by implementing :meth:`~SelfSupervisedRegressionTransformer._on_weights_updated` and :meth:`~SelfSupervisedRegressionTransformer._process`. **LRR.** For each channel *c*, predict it from the other channels in its -cluster via ridge regression, then subtract the prediction:: +group via ridge regression, then subtract the prediction:: y = X - X @ W = X @ (I - W) The effective weight matrix ``I - W`` is passed to :class:`~ezmsg.sigproc.affinetransform.AffineTransformTransformer`, which -automatically exploits block-diagonal structure when ``channel_clusters`` -are provided. +reads the block-diagonal structure off the weight matrix itself and picks +a dense or block matmul accordingly — the channel grouping is an input to +*fitting* only, never to applying. **Fitting.** Given data matrix *X* of shape ``(samples, channels)``, the sufficient statistic is the channel covariance ``C = X^T X``. When ``incremental=True`` (default), *C* is accumulated across :meth:`~SelfSupervisedRegressionTransformer.partial_fit` calls. -**Solving.** Within each cluster the weight matrix *W* is obtained from -the inverse of the (ridge-regularised) cluster covariance -``C_inv = (C_cluster + lambda * I)^{-1}`` using the block-inverse identity:: +**Solving.** Within each group the weight matrix *W* is obtained from +the inverse of the (ridge-regularised) group covariance +``C_inv = (C_group + lambda * I)^{-1}`` using the block-inverse identity:: W[:, c] = -C_inv[:, c] / C_inv[c, c], diag(W) = 0 This replaces the naive per-channel Cholesky loop with a single matrix -inverse per cluster, keeping the linear algebra in the source array +inverse per group, keeping the linear algebra in the source array namespace so that GPU-backed arrays benefit from device-side computation. """ @@ -58,19 +59,25 @@ AffineTransformTransformer, ) from ezmsg.sigproc.util.array import array_device, xp_create -from ezmsg.sigproc.util.channels import channel_clusters_from_field, validate_channel_clusters +from ezmsg.sigproc.util.channels import ( + ChannelGroupSpec, + group_spec_fields, + group_spec_fingerprint, + resolve_channel_groups, + validate_channel_groups, +) 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 -# regresses each channel against the *others* in its cluster, so a cluster with +# Minimum channels a group needs before it is rereferenced. Rereferencing +# regresses each channel against the *others* in its group, so a group 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 +# reference at all; 2 -> a single, degenerate reference). Such groups 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 +# group reduced to a channel or two (or an empty group) 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 +MIN_REREF_GROUP_SIZE = 3 # --------------------------------------------------------------------------- @@ -88,18 +95,19 @@ class SelfSupervisedRegressionSettings(ez.Settings): axis: str | None = None """Channel axis name. ``None`` defaults to the last dimension.""" - channel_clusters: list[list[int]] | None = None - """Per-cluster regression. ``None`` treats all channels as one cluster.""" + channel_groups: ChannelGroupSpec | None = None + """How to split the channel axis into groups for per-group regression: explicit + index groups (``[[0, 1, 2], [3, 4, 5]]``), the name of a structured field on the + channel coordinate axis (``"bank"`` to regress within each electrode bank), a + tuple of field names, or a callable. See + :data:`~ezmsg.sigproc.util.channels.ChannelGroupSpec`. - cluster_by_field: str | None = None - """Derive ``channel_clusters`` automatically from a structured field of the - channel coordinate axis (e.g. ``"bank"`` to regress within each electrode - bank). Used only when ``channel_clusters`` is None and the axis actually - carries that field; otherwise falls back to ``block_size`` (then a single - cluster). Explicit ``channel_clusters`` always takes precedence.""" + ``None`` -- or a field spec the incoming axis doesn't carry -- falls back to + ``block_size``, then to a single all-channel group.""" block_size: int | None = None - """If ``channel_clusters`` is ``None``, use this block size for an automatic clustering.""" + """Fallback grouping when ``channel_groups`` is ``None`` or resolves to nothing: + consecutive blocks of this many channels.""" ridge_lambda: float = 0.0 """Ridge (L2) regularisation parameter.""" @@ -114,9 +122,9 @@ class SelfSupervisedRegressionState: cxx: object | None = None # Array API; namespace matches source data. n_samples: int = 0 weights: object | None = None # Array API; namespace matches cxx. - resolved_clusters: list | None = None - """Clusters derived from ``cluster_by_field`` at reset (message-dependent), - cached so ``_get_channel_clusters`` can return them without the message.""" + resolved_groups: list | None = None + """``channel_groups`` resolved against the message at reset, cached so + ``_get_channel_groups`` can return them without one.""" class SelfSupervisedRegressionTransformer( @@ -138,37 +146,26 @@ class SelfSupervisedRegressionTransformer( def _hash_message(self, message: AxisArray) -> int: axis = self.settings.axis or message.dims[-1] axis_idx = message.get_axis_idx(axis) - components: tuple = (message.key, message.data.shape[axis_idx]) - # On the cluster_by_field path, re-derive clusters only when the channel - # axis gains or loses the target structured field -- a single O(1) - # presence boolean rather than O(channels) of field bytes hashed on every - # message (which grows with channel count on this hot path). Concession: - # if the axis is already structured and the channel count is unchanged, a - # change in the field's *values* (a live bank remap) is not detected. That - # is safe for real acquisition streams, whose channel->field map is static - # for the stream's life; a genuine remap arrives with a new key or a - # different channel count, both already folded in above. Mirrors the - # ezmsg-sigproc CommonRereference hash. - if self.settings.channel_clusters is None and self.settings.cluster_by_field is not None: - ax = message.axes.get(axis) - names = getattr(getattr(getattr(ax, "data", None), "dtype", None), "names", None) - components += (bool(names and self.settings.cluster_by_field in names),) - return hash(components) + # group_spec_fingerprint contributes an O(1) "can this spec resolve?" + # boolean rather than the field's bytes, so the per-message hash does not + # grow with channel count. See its docstring for what that deliberately + # does not detect. Mirrors the ezmsg-sigproc transformers' hash. + return hash( + (message.key, message.data.shape[axis_idx]) + + group_spec_fingerprint(message, axis, self.settings.channel_groups) + ) def _reset_state(self, message: AxisArray) -> None: axis = self.settings.axis or message.dims[-1] axis_idx = message.get_axis_idx(axis) n_channels = message.data.shape[axis_idx] - # Derive clusters from a structured channel-axis field (e.g. "bank") when - # requested and no explicit clusters were given. Cached so the - # message-less _get_channel_clusters can return them later. - if self.settings.channel_clusters is None and self.settings.cluster_by_field is not None: - self._state.resolved_clusters = channel_clusters_from_field(message, axis, self.settings.cluster_by_field) - else: - self._state.resolved_clusters = None + # Resolve the grouping against this message (a field- or callable-based + # spec needs one). Cached so the message-less _get_channel_groups can + # return it later. + self._state.resolved_groups = resolve_channel_groups(message, axis, self.settings.channel_groups) - self._validate_clusters(n_channels) + self._validate_groups(n_channels) self._state.cxx = None self._state.n_samples = 0 self._state.weights = None @@ -184,39 +181,51 @@ def _reset_state(self, message: AxisArray) -> None: self._state.weights = weights self._on_weights_updated() - # -- cluster validation -------------------------------------------------- - - def _get_channel_clusters(self, n_channels: int) -> list[list[int]] | None: - # Precedence: explicit channel_clusters > cluster_by_field-derived - # (cached at reset) > block_size > None (single cluster). - clusters = self.settings.channel_clusters - if clusters is None: - clusters = getattr(self._state, "resolved_clusters", None) - if clusters is None and self.settings.block_size is not None: - clusters = [ - list(range(i, min(i + self.settings.block_size, n_channels))) + # -- group resolution / validation --------------------------------------- + + def _static_channel_groups(self) -> list[np.ndarray] | None: + """Groups readable from the settings alone, without a message. + + Explicit index groups are static; field-name and callable specs depend on + the message and are resolved in :meth:`_reset_state` instead. This is what + lets the message-less :meth:`fit` honour explicit groups. + """ + spec = self.settings.channel_groups + if spec is None or callable(spec) or group_spec_fields(spec) is not None: + return None + return [np.asarray(group, dtype=np.intp).reshape(-1) for group in spec] + + def _get_channel_groups(self, n_channels: int) -> list[np.ndarray] | None: + # Precedence: resolved channel_groups (cached at reset, or static explicit + # groups) > block_size > None (single all-channel group). + groups = self._state.resolved_groups + if groups is None: + groups = self._static_channel_groups() + if groups is None and self.settings.block_size is not None: + groups = [ + np.arange(i, min(i + self.settings.block_size, n_channels), dtype=np.intp) for i in range(0, n_channels, self.settings.block_size) ] - return clusters - - 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 # implicit single cluster - if len(clusters) == 0: - # An empty cluster list is only legitimate with no channels (e.g. a + return groups + + def _validate_groups(self, n_channels: int) -> None: + """Raise if the resolved groups are empty, out of range, or overlapping.""" + groups = self._get_channel_groups(n_channels) + if groups is None: + return # implicit single group + if len(groups) == 0: + # An empty group 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 -- + # channel_groups=[], which would silently disable rereferencing -- # fail fast instead. 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." + f"channel_groups is empty but the input has {n_channels} channels. " + "Pass channel_groups=None to treat all channels as a single " + "group, or provide non-empty channel index groups." ) - validate_channel_clusters(clusters, n_channels) + validate_channel_groups(groups, n_channels) # -- weight solving ------------------------------------------------------ @@ -225,13 +234,13 @@ def _solve_weights(self, cxx): Uses the block-inverse identity: for target channel *c* with references *r*, ``w_c = -C_inv[r, c] / C_inv[c, c]`` where - ``C_inv = (C_cluster + λI)⁻¹``. This replaces the per-channel - Cholesky loop with one matrix inverse per cluster. + ``C_inv = (C_group + λI)⁻¹``. This replaces the per-channel + Cholesky loop with one matrix inverse per group. All computation stays in the source array namespace so that - GPU-backed arrays benefit from device-side execution. Cluster + GPU-backed arrays benefit from device-side execution. Group results are scattered into the full matrix via a selection-matrix - multiply (``S @ W_cluster @ S^T``) to avoid numpy fancy indexing. + multiply (``S @ W_group @ S^T``) to avoid numpy fancy indexing. Returns weight matrix *W* in the same namespace as *cxx*, with ``diag(W) == 0``. @@ -240,9 +249,9 @@ def _solve_weights(self, cxx): dev = array_device(cxx) n = cxx.shape[0] - clusters = self._get_channel_clusters(n) - if clusters is None: - clusters = [list(range(n))] + groups = self._get_channel_groups(n) + if groups is None: + groups = [np.arange(n, dtype=np.intp)] W = xp_create(xp.zeros, (n, n), dtype=cxx.dtype, device=dev) eye_n = xp_create(xp.eye, n, dtype=cxx.dtype, device=dev) @@ -251,24 +260,26 @@ def _solve_weights(self, cxx): # 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: + for group in groups: + idx = np.asarray(group, dtype=np.intp).reshape(-1) + k = idx.size + if k < MIN_REREF_GROUP_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. + # groups down to a single channel; never raises. continue - idx_xp = xp.asarray(cluster) if dev is None else xp.asarray(cluster, device=dev) + idx_list = idx.tolist() + idx_xp = xp.asarray(idx_list) if dev is None else xp.asarray(idx_list, device=dev) eye_k = xp_create(xp.eye, k, dtype=cxx.dtype, device=dev) - # Extract cluster sub-covariance (stays on device) + # Extract group sub-covariance (stays on device) sub = xp.take(xp.take(cxx, idx_xp, axis=0), idx_xp, axis=1) if self.settings.ridge_lambda > 0: sub = sub + self.settings.ridge_lambda * eye_k - # One inverse per cluster + # One inverse per group try: sub_inv = xp.linalg.inv(sub, **inv_kwargs) except Exception: @@ -278,18 +289,20 @@ def _solve_weights(self, cxx): diag_vals = xp.sum(sub_inv * eye_k, axis=0) # w_c = -C_inv[:, c] / C_inv[c, c], vectorised over all c - W_cluster = -(sub_inv / xp.reshape(diag_vals, (1, k))) + W_group = -(sub_inv / xp.reshape(diag_vals, (1, k))) # Zero the diagonal - W_cluster = W_cluster * (1.0 - eye_k) + W_group = W_group * (1.0 - eye_k) - # Scatter into full W - if k == n: - W = W + W_cluster + # Scatter into full W. The no-op shortcut needs the group to be every + # channel *in order* -- a callable spec may return all n permuted, and + # then the sub-block still has to be scattered back. + if k == n and np.array_equal(idx, np.arange(n, dtype=np.intp)): + W = W + W_group else: - # Selection matrix: columns of eye(n) at cluster indices + # Selection matrix: columns of eye(n) at group indices S = xp.take(eye_n, idx_xp, axis=1) # (n, k) - W = W + xp.matmul(S, xp.matmul(W_cluster, xp.permute_dims(S, (1, 0)))) + W = W + xp.matmul(S, xp.matmul(W_group, xp.permute_dims(S, (1, 0)))) return W @@ -341,7 +354,7 @@ def partial_fit(self, message: AxisArray) -> None: # type: ignore[override] 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) + self._validate_groups(n_channels) if n_channels == 0: # No channels to fit -- same 0-channel no-op as partial_fit. return @@ -374,15 +387,17 @@ def _process(self, message: AxisArray) -> AxisArray: ... class LRRSettings(SelfSupervisedRegressionSettings): """Settings for :class:`LRRTransformer`.""" - min_cluster_size: int = 32 - """Passed to :class:`AffineTransformTransformer` for the block-diagonal - merge threshold.""" + kernel: str = "auto" + """Forwarded to :attr:`~ezmsg.sigproc.affinetransform.AffineTransformSettings.kernel`. + ``"auto"`` lets the affine transformer choose between a dense and a + block-diagonal matmul from the structure of ``I - W``; ``"dense"`` / + ``"blocks"`` force it.""" 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 + yet. ``IDENTITY`` passes through (legacy); ``CAR`` applies per-group + leave-one-out common-average referencing from the resolved groups (groups + below :data:`MIN_REREF_GROUP_SIZE` stay identity, matching the fit's passthrough). Provided or fitted weights always take precedence over this cold-start default.""" @@ -390,6 +405,10 @@ class LRRSettings(SelfSupervisedRegressionSettings): @processor_state class LRRState(SelfSupervisedRegressionState): affine: AffineTransformTransformer | None = None + effective: object | None = None + """Latest ``I - W``, in the namespace of the fitted weights. Held here rather + than pushed straight into an affine transformer because the affine is not + built until a message actually needs it -- see :meth:`LRRTransformer._process`.""" class LRRTransformer( @@ -405,35 +424,44 @@ class LRRTransformer( def _reset_state(self, message: AxisArray) -> None: self._state.affine = None + self._state.effective = None super()._reset_state(message) # -- weights → affine transform ----------------------------------------- + def _make_affine(self, effective) -> AffineTransformTransformer: + # No channel_groups: the affine derives block structure from the weight + # matrix itself, and grouping only ever builds kind/callable weights -- + # which these are not. + return AffineTransformTransformer( + AffineTransformSettings( + weights=effective, + axis=self.settings.axis, + kernel=self.settings.kernel, + ) + ) + def _on_weights_updated(self) -> None: xp = get_namespace(self._state.weights) dev = array_device(self._state.weights) n = self._state.weights.shape[0] effective = xp_create(xp.eye, n, dtype=self._state.weights.dtype, device=dev) - self._state.weights - - # Prefer in-place weight update when the affine transformer supports - # it (avoids a full _reset_state round-trip on every partial_fit). + self._state.effective = effective + + # Update an existing affine in place (avoids a full _reset_state + # round-trip on every partial_fit). The default recalc_structure=False is + # what we want: refitting changes the weight *values*, not their sparsity + # pattern, which is fixed by the channel grouping. + # + # Do NOT construct the affine here when there isn't one. An affine built + # now would carry these weights in its *settings*, and its first + # _reset_state -- which does not happen until a message arrives -- rebuilds + # its state from those settings. Any refit in between would update state + # that is about to be overwritten, so a stream that fits several times + # before its first signal message would silently apply the *first* fit + # forever. _process builds it instead, from the latest weights. 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=None, - min_cluster_size=self.settings.min_cluster_size, - ) - ) # -- transform ----------------------------------------------------------- @@ -442,32 +470,30 @@ def _process(self, message: AxisArray) -> AxisArray: 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. + # from empty channel groups would raise downstream. return message if self._state.affine is None: - axis_idx = message.get_axis_idx(axis) - n_channels = message.data.shape[axis_idx] - - # 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, - axis=self.settings.axis, - channel_clusters=self._get_channel_clusters(n_channels), - min_cluster_size=self.settings.min_cluster_size, + effective = self._state.effective + if effective is None: + axis_idx = message.get_axis_idx(axis) + n_channels = message.data.shape[axis_idx] + + # No weights provided or fit yet: build the configured cold-start + # default (identity, or per-group leave-one-out CAR matching the + # fit's passthrough for groups below MIN_REREF_GROUP_SIZE). + # Built as numpy; the affine transformer converts weights to the + # message's namespace/dtype/device on first use. + groups = self._get_channel_groups(n_channels) + effective = rereference_matrix( + self.settings.init_default, + n_channels, + groups=None if groups is None else [group.tolist() for group in groups], + include_current=False, + min_reref_size=MIN_REREF_GROUP_SIZE, ) - ) + # Deferred to here so the affine is built from the newest weights: any + # number of partial_fit calls may have landed since the last message. + self._state.affine = self._make_affine(effective) return self._state.affine(message) diff --git a/tests/benchmark/bench_lrr.py b/tests/benchmark/bench_lrr.py index d77955c..ca02614 100644 --- a/tests/benchmark/bench_lrr.py +++ b/tests/benchmark/bench_lrr.py @@ -22,14 +22,14 @@ # --------------------------------------------------------------------------- N_CH = 512 -N_CLUSTERS = 8 -CLUSTER_SIZE = N_CH // N_CLUSTERS # 64 +N_GROUPS = 8 +GROUP_SIZE = N_CH // N_GROUPS # 64 FS = 30_000.0 CHUNK_SIZES = [20, 50, 100, 150, 200, 300] WARMUP_ITERS = 20 BENCH_ITERS = 200 -CLUSTERS = [list(range(i * CLUSTER_SIZE, (i + 1) * CLUSTER_SIZE)) for i in range(N_CLUSTERS)] +GROUPS = [list(range(i * GROUP_SIZE, (i + 1) * GROUP_SIZE)) for i in range(N_GROUPS)] # --------------------------------------------------------------------------- @@ -89,14 +89,14 @@ def _bench_loop_sync(fn, sync_fn, n_warmup: int, n_iters: int) -> list[float]: def bench_process_numpy() -> None: _print_header("_process (inference) — NumPy") - print(f" {N_CH} channels, {N_CLUSTERS}x{CLUSTER_SIZE} clusters, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") + print(f" {N_CH} channels, {N_GROUPS}x{GROUP_SIZE} groups, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") print() rng = np.random.default_rng(0) # Fit via partial_fit so the message hash is primed for send() fit_data = rng.standard_normal((2000, N_CH)) - proc = LRRTransformer(LRRSettings(channel_clusters=CLUSTERS, min_cluster_size=1)) + proc = LRRTransformer(LRRSettings(channel_groups=GROUPS)) proc.partial_fit(_make_msg(fit_data)) for chunk in CHUNK_SIZES: @@ -113,11 +113,11 @@ def bench_process_numpy() -> None: def bench_partial_fit_numpy() -> None: _print_header("partial_fit (training) — NumPy") - print(f" {N_CH} channels, {N_CLUSTERS}x{CLUSTER_SIZE} clusters, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") + print(f" {N_CH} channels, {N_GROUPS}x{GROUP_SIZE} groups, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") print() rng = np.random.default_rng(1) - proc = LRRTransformer(LRRSettings(channel_clusters=CLUSTERS, min_cluster_size=1)) + proc = LRRTransformer(LRRSettings(channel_groups=GROUPS)) for chunk in CHUNK_SIZES: data = rng.standard_normal((chunk, N_CH)) @@ -144,7 +144,7 @@ def bench_process_mps() -> None: return _print_header("_process (inference) — Torch MPS") - print(f" {N_CH} channels, {N_CLUSTERS}x{CLUSTER_SIZE} clusters, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") + print(f" {N_CH} channels, {N_GROUPS}x{GROUP_SIZE} groups, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") print() rng = np.random.default_rng(0) @@ -152,7 +152,7 @@ def bench_process_mps() -> None: # Fit on CPU (numpy), then send MPS data to trigger device conversion fit_data = rng.standard_normal((2000, N_CH)) - proc = LRRTransformer(LRRSettings(channel_clusters=CLUSTERS, min_cluster_size=1)) + proc = LRRTransformer(LRRSettings(channel_groups=GROUPS)) proc.partial_fit(_make_msg(fit_data)) def sync(): @@ -178,12 +178,12 @@ def bench_partial_fit_mps() -> None: return _print_header("partial_fit (training) — Torch MPS") - print(f" {N_CH} channels, {N_CLUSTERS}x{CLUSTER_SIZE} clusters, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") + print(f" {N_CH} channels, {N_GROUPS}x{GROUP_SIZE} groups, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") print() _ = np.random.default_rng(1) device = torch.device("mps") - proc = LRRTransformer(LRRSettings(channel_clusters=CLUSTERS, min_cluster_size=1)) + proc = LRRTransformer(LRRSettings(channel_groups=GROUPS)) def sync(): torch.mps.synchronize() @@ -213,14 +213,14 @@ def bench_process_mlx() -> None: return _print_header("_process (inference) — MLX") - print(f" {N_CH} channels, {N_CLUSTERS}x{CLUSTER_SIZE} clusters, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") + print(f" {N_CH} channels, {N_GROUPS}x{GROUP_SIZE} groups, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") print() rng = np.random.default_rng(0) # Fit on CPU (numpy), then send MLX data fit_data = rng.standard_normal((2000, N_CH)) - proc = LRRTransformer(LRRSettings(channel_clusters=CLUSTERS, min_cluster_size=1)) + proc = LRRTransformer(LRRSettings(channel_groups=GROUPS)) proc.partial_fit(_make_msg(fit_data)) def sync(): @@ -251,7 +251,7 @@ def bench_partial_fit_mlx() -> None: return _print_header("partial_fit (training) — MLX") - print(f" {N_CH} channels, {N_CLUSTERS}x{CLUSTER_SIZE} clusters, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") + print(f" {N_CH} channels, {N_GROUPS}x{GROUP_SIZE} groups, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") # MLX linalg.inv doesn't support GPU yet; run inv on CPU stream print(" NOTE: linalg.inv runs on mx.cpu stream (GPU not supported)") print() @@ -259,7 +259,7 @@ def bench_partial_fit_mlx() -> None: import mlx.core as mx _ = np.random.default_rng(1) - proc = LRRTransformer(LRRSettings(channel_clusters=CLUSTERS, min_cluster_size=1)) + proc = LRRTransformer(LRRSettings(channel_groups=GROUPS)) # Monkey-patch _solve_weights to use mx.cpu stream for inv original_solve = proc._solve_weights @@ -301,7 +301,7 @@ def run(): # --------------------------------------------------------------------------- if __name__ == "__main__": - print(f"LRRTransformer benchmark: {N_CH} channels, {N_CLUSTERS} clusters of {CLUSTER_SIZE}, fs={FS / 1e3:.0f} kHz") + print(f"LRRTransformer benchmark: {N_CH} channels, {N_GROUPS} groups of {GROUP_SIZE}, fs={FS / 1e3:.0f} kHz") bench_process_numpy() bench_partial_fit_numpy() diff --git a/tests/unit/test_adaptive_linear_regressor.py b/tests/unit/test_adaptive_linear_regressor.py index 6d59499..d8d6274 100644 --- a/tests/unit/test_adaptive_linear_regressor.py +++ b/tests/unit/test_adaptive_linear_regressor.py @@ -1,9 +1,8 @@ import numpy as np import pytest import sklearn.linear_model -from ezmsg.sigproc.window import WindowTransformer -from ezmsg.sigproc.window import WindowSettings from ezmsg.baseproc import SampleTriggerMessage +from ezmsg.sigproc.window import WindowSettings, WindowTransformer from ezmsg.util.messages.axisarray import AxisArray, replace from ezmsg.learn.process.adaptive_linear_regressor import AdaptiveLinearRegressorTransformer diff --git a/tests/unit/test_ssr.py b/tests/unit/test_ssr.py index 152b64f..83ffe05 100644 --- a/tests/unit/test_ssr.py +++ b/tests/unit/test_ssr.py @@ -7,7 +7,7 @@ from ezmsg.util.messages.axisarray import AxisArray from ezmsg.learn.process.ssr import ( - MIN_REREF_CLUSTER_SIZE, + MIN_REREF_GROUP_SIZE, LRRSettings, LRRTransformer, RereferenceKind, @@ -111,33 +111,33 @@ def test_diagonal_zero(self): np.testing.assert_array_equal(np.diag(proc.state.weights), 0.0) -class TestChannelClusters: - def test_channel_clusters(self): - """Cross-cluster weights must be zero; within-cluster weights non-zero.""" +class TestChannelGroups: + def test_channel_groups(self): + """Cross-group weights must be zero; within-group weights non-zero.""" rng = np.random.default_rng(3) n_ch = 8 - clusters = [[0, 1, 2, 3], [4, 5, 6, 7]] + groups = [[0, 1, 2, 3], [4, 5, 6, 7]] X = _random_data(n_ch=n_ch, rng=rng) msg = _make_axisarray(X) - proc = LRRTransformer(LRRSettings(channel_clusters=clusters)) + proc = LRRTransformer(LRRSettings(channel_groups=groups)) proc.partial_fit(msg) W = proc.state.weights - # Cross-cluster should be zero - for c1 in clusters: - for c2 in clusters: + # Cross-group should be zero + for c1 in groups: + for c2 in groups: if c1 is c2: continue cross = W[np.ix_(c1, c2)] np.testing.assert_array_equal(cross, 0.0) - # Within-cluster (off-diagonal) should be non-zero - for cluster in clusters: - sub = W[np.ix_(cluster, cluster)] - off_diag = sub[~np.eye(len(cluster), dtype=bool)] - assert np.any(off_diag != 0), "Expected non-zero within-cluster weights" + # Within-group (off-diagonal) should be non-zero + for group in groups: + sub = W[np.ix_(group, group)] + off_diag = sub[~np.eye(len(group), dtype=bool)] + assert np.any(off_diag != 0), "Expected non-zero within-group weights" def _banked_axisarray(data: np.ndarray, banks: list[str], key: str = "test") -> AxisArray: @@ -156,18 +156,18 @@ def _banked_axisarray(data: np.ndarray, banks: list[str], key: str = "test") -> ) -class TestClusterByField: - def test_bank_field_matches_explicit_clusters(self): - """cluster_by_field='bank' derives the same clusters (and weights) as - passing the equivalent explicit channel_clusters.""" +class TestGroupByField: + def test_bank_field_matches_explicit_groups(self): + """channel_groups='bank' derives the same groups (and weights) as + passing the equivalent explicit channel_groups.""" rng = np.random.default_rng(7) banks = ["A", "A", "A", "A", "B", "B", "B", "B"] X = _random_data(n_ch=len(banks), rng=rng) - proc_field = LRRTransformer(LRRSettings(axis="ch", cluster_by_field="bank")) + proc_field = LRRTransformer(LRRSettings(axis="ch", channel_groups="bank")) proc_field.partial_fit(_banked_axisarray(X, banks)) - proc_explicit = LRRTransformer(LRRSettings(axis="ch", channel_clusters=[[0, 1, 2, 3], [4, 5, 6, 7]])) + proc_explicit = LRRTransformer(LRRSettings(axis="ch", channel_groups=[[0, 1, 2, 3], [4, 5, 6, 7]])) proc_explicit.partial_fit(_make_axisarray(X)) np.testing.assert_array_equal(proc_field.state.weights, proc_explicit.state.weights) @@ -175,26 +175,52 @@ def test_bank_field_matches_explicit_clusters(self): W = proc_field.state.weights np.testing.assert_array_equal(W[np.ix_([0, 1, 2, 3], [4, 5, 6, 7])], 0.0) - def test_explicit_clusters_take_precedence(self): - """Explicit channel_clusters win over cluster_by_field.""" + def test_callable_spec(self): + """A callable spec is resolved against the message like any other.""" rng = np.random.default_rng(8) banks = ["A", "A", "A", "A", "B", "B", "B", "B"] X = _random_data(n_ch=len(banks), rng=rng) - # One all-channel cluster should override the bank grouping. - proc = LRRTransformer(LRRSettings(axis="ch", channel_clusters=[list(range(8))], cluster_by_field="bank")) + # A callable returning one all-channel group: cross-"bank" weights are + # then NOT forced to zero. + proc = LRRTransformer(LRRSettings(axis="ch", channel_groups=lambda msg, axis: [list(range(8))])) proc.partial_fit(_banked_axisarray(X, banks)) - # With a single cluster, cross-"bank" weights are NOT forced to zero. W = proc.state.weights assert np.any(W[np.ix_([0, 1, 2, 3], [4, 5, 6, 7])] != 0) + def test_multi_field_spec(self): + """A tuple of field names groups by their combination.""" + n_ch = 8 + dt = np.dtype([("array", "U1"), ("bank", "U1")]) + ch = np.zeros(n_ch, dtype=dt) + ch["array"] = ["1", "1", "1", "1", "2", "2", "2", "2"] + ch["bank"] = ["A", "A", "B", "B", "A", "A", "B", "B"] + X = _random_data(n_ch=n_ch, rng=np.random.default_rng(21)) + 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", channel_groups=("array", "bank"))) + proc.partial_fit(msg) + + # (array, bank) -> four groups of 2, each below MIN_REREF_GROUP_SIZE, so + # every weight stays zero (identity passthrough). + assert [g.tolist() for g in proc.state.resolved_groups] == [[0, 1], [2, 3], [4, 5], [6, 7]] + np.testing.assert_array_equal(proc.state.weights, 0.0) + def test_missing_field_falls_back_to_block_size(self): - """cluster_by_field with no structured bank field falls back to block_size.""" + """channel_groups with no structured bank field falls back to block_size.""" rng = np.random.default_rng(9) n_ch = 8 X = _random_data(n_ch=n_ch, rng=rng) # Plain axis (no structured bank field) + block_size=4 -> two contiguous blocks. - proc_field = LRRTransformer(LRRSettings(axis="ch", cluster_by_field="bank", block_size=4)) + proc_field = LRRTransformer(LRRSettings(axis="ch", channel_groups="bank", block_size=4)) proc_field.partial_fit(_make_axisarray(X)) proc_block = LRRTransformer(LRRSettings(axis="ch", block_size=4)) @@ -210,21 +236,21 @@ def test_bank_field_value_change_is_not_detected(self): real hardware arrives with a new key or channel count (escape hatch below).""" rng = np.random.default_rng(11) X = _random_data(n_ch=4, rng=rng) - proc = LRRTransformer(LRRSettings(axis="ch", cluster_by_field="bank")) + proc = LRRTransformer(LRRSettings(axis="ch", channel_groups="bank")) - # First arrangement: banks A,A,B,B -> clusters {0,1},{2,3}. + # First arrangement: banks A,A,B,B -> groups {0,1},{2,3}. proc.partial_fit(_banked_axisarray(X, ["A", "A", "B", "B"], key="x")) - assert proc.state.resolved_clusters == [[0, 1], [2, 3]] + assert [g.tolist() for g in proc.state.resolved_groups] == [[0, 1], [2, 3]] np.testing.assert_array_equal(proc.state.weights[np.ix_([0, 1], [2, 3])], 0.0) # Same key + channel count, different banks -> hash unchanged, so the - # cached clusters are (deliberately) NOT re-derived. + # cached groups are (deliberately) NOT re-derived. proc.partial_fit(_banked_axisarray(X, ["A", "B", "A", "B"], key="x")) - assert proc.state.resolved_clusters == [[0, 1], [2, 3]] + assert [g.tolist() for g in proc.state.resolved_groups] == [[0, 1], [2, 3]] # Escape hatch: a new key (as a real remap would carry) forces re-derivation. proc.partial_fit(_banked_axisarray(X, ["A", "B", "A", "B"], key="y")) - assert proc.state.resolved_clusters == [[0, 2], [1, 3]] + assert [g.tolist() for g in proc.state.resolved_groups] == [[0, 2], [1, 3]] class TestIncrementalAccumulates: @@ -354,6 +380,51 @@ def test_partial_fit_transform(self): np.testing.assert_allclose(out1.data, out2.data, atol=1e-12) +class TestRefitBeforeFirstMessage: + """Regression: the internal affine must be built from the NEWEST weights. + + ``LRRUnit`` takes training on ``INPUT_SAMPLE`` and signal on a separate + stream, so several ``partial_fit`` calls routinely land before the first + message is processed. Building the affine eagerly in ``_on_weights_updated`` + put the first fit's weights in the affine's *settings*; its ``_reset_state`` + (deferred until a message arrives) then rebuilt from those settings and + discarded every later fit, silently applying the first fit forever. + """ + + def test_multiple_fits_before_first_message(self): + rng = np.random.default_rng(0) + X1 = _random_data(n_ch=4, rng=rng) + X2 = _random_data(n_ch=4, rng=rng) + + proc = LRRTransformer(LRRSettings(incremental=False)) + proc.partial_fit(_make_axisarray(X1)) + proc.partial_fit(_make_axisarray(X2)) # no message processed in between + out = proc.send(_make_axisarray(X2)) + + expected = X2 @ (np.eye(4) - proc.state.weights) + np.testing.assert_allclose(out.data, expected, atol=1e-10) + + # And it is NOT the stale first fit. + stale = LRRTransformer(LRRSettings(incremental=False)) + stale.partial_fit(_make_axisarray(X1)) + assert not np.allclose(out.data, X2 @ (np.eye(4) - stale.state.weights), atol=1e-8) + + def test_refit_after_first_message_still_applies(self): + """The in-place set_weights path (affine already built) stays correct.""" + rng = np.random.default_rng(1) + X1 = _random_data(n_ch=4, rng=rng) + X2 = _random_data(n_ch=4, rng=rng) + + proc = LRRTransformer(LRRSettings(incremental=False)) + proc.partial_fit(_make_axisarray(X1)) + proc.send(_make_axisarray(X1)) # builds the affine + proc.partial_fit(_make_axisarray(X2)) + out = proc.send(_make_axisarray(X2)) + + expected = X2 @ (np.eye(4) - proc.state.weights) + np.testing.assert_allclose(out.data, expected, atol=1e-10) + + class TestPassthroughThenFit: def test_passthrough_then_fit(self): """Pre-fit send() should passthrough, then partial_fit() should update weights.""" @@ -374,37 +445,58 @@ def test_passthrough_then_fit(self): np.testing.assert_allclose(out_after.data, expected, atol=1e-10) -class TestInvalidClusterIndicesRaise: - def test_invalid_cluster_indices_raise(self): - """Out-of-range indices in channel_clusters should raise ValueError.""" +class TestInvalidGroupIndicesRaise: + def test_invalid_group_indices_raise(self): + """Out-of-range indices in channel_groups should raise ValueError.""" rng = np.random.default_rng(11) X = _random_data(n_ch=4, rng=rng) msg = _make_axisarray(X) - proc = LRRTransformer(LRRSettings(channel_clusters=[[0, 1, 99]])) + proc = LRRTransformer(LRRSettings(channel_groups=[[0, 1, 99]])) with pytest.raises(ValueError, match="out-of-range"): proc.partial_fit(msg) -class TestClustersEngageBlockDiagonal: - def test_clusters_engage_block_diagonal(self): - """When clusters create a block-diagonal I-W, AffineTransform uses cluster opt.""" +class TestGroupsEngageBlockDiagonal: + def test_groups_engage_block_diagonal(self): + """kernel='blocks' forces the block-diagonal matmul over the blocks the + affine reads off I - W; the result must equal the dense matmul.""" rng = np.random.default_rng(12) n_ch = 8 - clusters = [[0, 1, 2, 3], [4, 5, 6, 7]] + groups = [[0, 1, 2, 3], [4, 5, 6, 7]] X = _random_data(n_ch=n_ch, n_times=300, rng=rng) msg = _make_axisarray(X) - proc = LRRTransformer(LRRSettings(channel_clusters=clusters, min_cluster_size=1)) + proc = LRRTransformer(LRRSettings(channel_groups=groups, kernel="blocks")) proc.partial_fit(msg) out = proc.send(msg) + # The affine really is on the block path, and it found the two 4-ch blocks. + assert proc.state.affine.state.blocks is not None + assert len(proc.state.affine.state.blocks) == 2 + # Verify output is correct — the block-diagonal path should produce # the same result as a full matmul. W = proc.state.weights expected = X @ (np.eye(n_ch) - W) np.testing.assert_allclose(out.data, expected, atol=1e-10) + def test_kernel_dense_forces_dense(self): + """kernel='dense' keeps the full matrix even when I - W is block-diagonal.""" + rng = np.random.default_rng(13) + n_ch = 8 + X = _random_data(n_ch=n_ch, n_times=300, rng=rng) + msg = _make_axisarray(X) + + proc = LRRTransformer(LRRSettings(channel_groups=[[0, 1, 2, 3], [4, 5, 6, 7]], kernel="dense")) + proc.partial_fit(msg) + out = proc.send(msg) + + assert proc.state.affine.state.blocks is None + assert proc.state.affine.state.weights is not None + expected = X @ (np.eye(n_ch) - proc.state.weights) + np.testing.assert_allclose(out.data, expected, atol=1e-10) + class TestPrecalculatedWeights: def test_precalculated_weights(self): @@ -453,19 +545,22 @@ def test_precalculated_weights_from_file(self): class TestApplyFollowsWeightBlocks: """Regression: applying fit/loaded weights must equal ``X @ (I - W)`` no matter - what ``block_size`` / ``cluster_by_field`` the transformer carries. + what ``block_size`` / ``channel_groups`` 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 + block structure, not a group hint that may not match it. Previously the + affine was built with ``_get_channel_groups()``, which falls back to + ``block_size`` when the ``channel_groups`` 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 -- + field). When those fallback groups 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. + + ezmsg-sigproc#198 removed the hint's ability to matter at all: structure is + now always read off W. This test keeps the guarantee pinned from this side. """ def test_apply_ignores_block_size_finer_than_weight_blocks(self): @@ -476,7 +571,7 @@ def test_apply_ignores_block_size_finer_than_weight_blocks(self): 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 = LRRTransformer(LRRSettings(channel_groups=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. @@ -486,30 +581,31 @@ def test_apply_ignores_block_size_finer_than_weight_blocks(self): 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 + # Apply the loaded W with block_size=32 (finer than the W's 64-ch blocks). + # These 4x32 groups 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)) + # kernel="blocks" forces the block path so the guarantee is tested there + # rather than at whatever the planner happens to pick for this size. + proc = LRRTransformer(LRRSettings(weights=W, block_size=32, kernel="blocks")) 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: - """Clusters smaller than MIN_REREF_CLUSTER_SIZE (and empty inputs) pass + """Groups smaller than MIN_REREF_GROUP_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")) + proc = LRRTransformer(LRRSettings(axis="ch", channel_groups="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")) + proc = LRRTransformer(LRRSettings(axis="ch", channel_groups="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 [] @@ -529,25 +625,25 @@ def test_single_channel_identity(self): 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 + """A group with < MIN_REREF_GROUP_SIZE channels is left untouched.""" + n = MIN_REREF_GROUP_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 + """A group with exactly MIN_REREF_GROUP_SIZE channels is rereferenced.""" + n = MIN_REREF_GROUP_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 + def test_mixed_small_and_large_groups(self): + """Per-group: a full bank rereferences while a lone-channel bank in the same message passes through untouched.""" - big = MIN_REREF_CLUSTER_SIZE + 1 + big = MIN_REREF_GROUP_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 @@ -555,49 +651,49 @@ def test_mixed_small_and_large_clusters(self): 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 + def test_empty_explicit_groups_with_channels_raises(self): + """channel_groups=[] 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=[])) + proc = LRRTransformer(LRRSettings(axis="ch", channel_groups=[])) 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 + """init_default=CAR: cold-start per-group 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.""" + def _loo_car(X: np.ndarray, groups) -> np.ndarray: + """Reference per-group 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: + for cl in groups: + if len(cl) < MIN_REREF_GROUP_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]] + def test_car_applies_leave_one_out_per_group(self): + groups = [[0, 1, 2, 3], [4, 5, 6, 7]] X = _random_data(n_ch=8) - proc = LRRTransformer(LRRSettings(channel_clusters=clusters, init_default=RereferenceKind.CAR)) + proc = LRRTransformer(LRRSettings(channel_groups=groups, 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) + np.testing.assert_allclose(out.data, self._loo_car(X, groups), 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]] + def test_car_leaves_small_groups_identity(self): + # first group (size 2 < MIN_REREF_GROUP_SIZE) must pass through + groups = [[0, 1], [2, 3, 4, 5, 6, 7]] X = _random_data(n_ch=8) - proc = LRRTransformer(LRRSettings(channel_clusters=clusters, init_default=RereferenceKind.CAR)) + proc = LRRTransformer(LRRSettings(channel_groups=groups, 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) + np.testing.assert_allclose(out.data, self._loo_car(X, groups), atol=1e-10) def test_car_from_bank_field(self): - """cluster_by_field='bank' + CAR reproduces per-bank leave-one-out CAR.""" + """channel_groups='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" @@ -611,14 +707,14 @@ def test_car_from_bank_field(self): }, key="test", ) - proc = LRRTransformer(LRRSettings(axis="ch", cluster_by_field="bank", init_default=RereferenceKind.CAR)) + proc = LRRTransformer(LRRSettings(axis="ch", channel_groups="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]])) + proc = LRRTransformer(LRRSettings(channel_groups=[[0, 1, 2, 3], [4, 5, 6, 7]])) out = proc.send(_make_axisarray(X)) np.testing.assert_allclose(out.data, X, atol=1e-12) @@ -633,17 +729,17 @@ def test_provided_weights_override_car(self): 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]] + groups = [[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 = LRRTransformer(LRRSettings(channel_groups=groups, 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) + assert not np.allclose(out.data, self._loo_car(X, groups), atol=1e-8) # --------------------------------------------------------------------------- @@ -669,19 +765,19 @@ class TestBackendPreservation: 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]] + GROUPS = [[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).""" + """All weight arrays held by the internal affine (dense or per-group).""" if affine.state.weights is not None: return [affine.state.weights] - return [sub_w for _, _, sub_w in affine.state.clusters] + return [sub_w for _, _, sub_w in affine.state.blocks] 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)) + proc = LRRTransformer(LRRSettings(channel_groups=self.GROUPS, init_default=RereferenceKind.CAR)) out = proc.send(_make_axisarray(conv(X.copy()))) assert isinstance(out.data, typ) @@ -691,7 +787,7 @@ def test_cold_start_car_converts_and_preserves(self, backend): assert isinstance(w, typ) # Values match the numpy cold-start CAR. - ref_proc = LRRTransformer(LRRSettings(channel_clusters=self.CLUSTERS, init_default=RereferenceKind.CAR)) + ref_proc = LRRTransformer(LRRSettings(channel_groups=self.GROUPS, init_default=RereferenceKind.CAR)) ref = ref_proc.send(_make_axisarray(X)) np.testing.assert_allclose(np.asarray(out.data), ref.data, atol=1e-5) @@ -699,7 +795,7 @@ 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 = LRRTransformer(LRRSettings(channel_groups=self.GROUPS)) proc.partial_fit(msg) assert isinstance(proc.state.cxx, typ) @@ -711,7 +807,7 @@ def test_fit_keeps_state_and_output_in_backend(self, backend): assert isinstance(w, typ) # Fitted output matches the numpy fit within float32 tolerance. - ref_proc = LRRTransformer(LRRSettings(channel_clusters=self.CLUSTERS)) + ref_proc = LRRTransformer(LRRSettings(channel_groups=self.GROUPS)) 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) @@ -720,12 +816,12 @@ 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 = LRRTransformer(LRRSettings(channel_groups=self.GROUPS)) 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)) + proc = LRRTransformer(LRRSettings(weights=W, channel_groups=self.GROUPS)) out = proc.send(_make_axisarray(conv(X.copy()))) assert isinstance(out.data, typ) for w in self._affine_weight_arrays(proc.state.affine):