ezmsg.learn.process.ssr#

Self-supervised regression framework and LRR implementation.

This module provides a general framework for self-supervised channel regression via SelfSupervisedRegressionTransformer, and a concrete implementation — Linear Regression Rereferencing (LRR) — via LRRTransformer.

Framework. The base class accumulates the channel covariance 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 _on_weights_updated() and _process().

LRR. For each channel c, predict it from the other channels in its group via ridge regression, then subtract the prediction:

y = X - X @ W = X @ (I - W)

The effective weight matrix I - W is passed to AffineTransformTransformer, which 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 partial_fit() calls.

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 group, keeping the linear algebra in the source array namespace so that GPU-backed arrays benefit from device-side computation.

Classes

class LRRSettings(weights=None, axis=None, channel_groups=None, block_size=None, ridge_lambda=0.0, incremental=True, kernel='auto', init_default=RereferenceKind.IDENTITY)[source]#

Bases: SelfSupervisedRegressionSettings

Settings for LRRTransformer.

Parameters:

  • weights (np.ndarray | str | Path | None)

  • axis (str | None)

  • channel_groups (ChannelGroupSpec | None)

  • block_size (int | None)

  • ridge_lambda (float)

  • incremental (bool)

  • kernel (str)

  • init_default (RereferenceKind)

kernel: str = 'auto'#

Forwarded to 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 = 'identity'#

Effective transform used when weights is None and nothing has been fit yet. IDENTITY passes through (legacy); CAR applies per-group leave-one-out common-average referencing from the resolved groups (groups below MIN_REREF_GROUP_SIZE stay identity, matching the fit’s passthrough). Provided or fitted weights always take precedence over this cold-start default.

__init__(weights=None, axis=None, channel_groups=None, block_size=None, ridge_lambda=0.0, incremental=True, kernel='auto', init_default=RereferenceKind.IDENTITY)#
Parameters:

  • weights (np.ndarray | str | Path | None)

  • axis (str | None)

  • channel_groups (ChannelGroupSpec | None)

  • block_size (int | None)

  • ridge_lambda (float)

  • incremental (bool)

  • kernel (str)

  • init_default (RereferenceKind)

Return type:

None

class LRRState[source]#

Bases: 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 LRRTransformer._process().

class LRRTransformer(*args, **kwargs)[source]#

Bases: SelfSupervisedRegressionTransformer[LRRSettings, LRRState]

Adaptive LRR transformer.

partial_fit accepts a plain AxisArray (self-supervised), and the transform step is delegated to an internal AffineTransformTransformer.

class LRRUnit(*args, settings=None, **kwargs)[source]#

Bases: BaseAdaptiveTransformerUnit[LRRSettings, AxisArray, AxisArray, LRRTransformer]

ezmsg Unit wrapping LRRTransformer.

Follows the BaseAdaptiveDecompUnit pattern — accepts AxisArray for self-supervised training via INPUT_SAMPLE.

Parameters:

settings (Settings | None)

SETTINGS#

alias of LRRSettings

INPUT_SAMPLE = InputStream:unlocated[AxisArray]()#
async on_sample(msg)[source]#
Return type:

None

Parameters:

msg (AxisArray)

class SelfSupervisedRegressionSettings(weights=None, axis=None, channel_groups=None, block_size=None, ridge_lambda=0.0, incremental=True)[source]#

Bases: Settings

Settings common to all self-supervised regression modes.

Parameters:

  • weights (np.ndarray | str | Path | None)

  • axis (str | None)

  • channel_groups (ChannelGroupSpec | None)

  • block_size (int | None)

  • ridge_lambda (float)

  • incremental (bool)

weights: np.ndarray | str | Path | None = None#

Pre-calculated weight matrix W or path to a CSV file (np.loadtxt compatible). If provided, the transformer is ready immediately.

axis: str | None = None#

Channel axis name. None defaults to the last dimension.

channel_groups: ChannelGroupSpec | None = None#

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 ChannelGroupSpec.

None – or a field spec the incoming axis doesn’t carry – falls back to block_size, then to a single all-channel group.

Type:

How to split the channel axis into groups for per-group regression

block_size: int | None = None#

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.

incremental: bool = True#

When True, accumulate X^T X across partial_fit() calls. When False, each call replaces the previous statistics.

__init__(weights=None, axis=None, channel_groups=None, block_size=None, ridge_lambda=0.0, incremental=True)#
Parameters:

  • weights (np.ndarray | str | Path | None)

  • axis (str | None)

  • channel_groups (ChannelGroupSpec | None)

  • block_size (int | None)

  • ridge_lambda (float)

  • incremental (bool)

Return type:

None

class SelfSupervisedRegressionState[source]#

Bases: object

cxx: object | None = None#
n_samples: int = 0#
weights: object | None = None#
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(*args, **kwargs)[source]#

Bases: BaseAdaptiveTransformer[SettingsType, AxisArray, AxisArray, StateType], Generic[SettingsType, StateType]

Abstract base for self-supervised regression transformers.

Subclasses must implement:

  • _on_weights_updated() — called whenever the weight matrix W is (re)computed, so the subclass can build whatever internal transform it needs (e.g. I - W for LRR).

  • _process() — the per-message transform step.

partial_fit(message)[source]#
Return type:

None

Parameters:

message (AxisArray)

fit(X)[source]#

Batch fit from a raw numpy array (samples x channels).

Return type:

None

Parameters:

X (ndarray)