From b24561e9be5ab864bc6d00ff56f57dcd6b1a0b51 Mon Sep 17 00:00:00 2001 From: Hamza Abdelhedi Date: Fri, 28 Aug 2026 15:43:55 +0200 Subject: [PATCH] MAINT: Delegate SSP-SIR reconstruction to MNE --- docs/api.rst | 2 - docs/changes/devel/90.feature.rst | 2 + examples/sspsir/plot_01_sspsir_basics.py | 10 +- mne_denoise/_leadfield.py | 271 +++++++++++++---------- mne_denoise/sspsir.py | 200 ++++++----------- tests/test_leadfield.py | 53 ++++- tests/test_mne.py | 4 - tests/test_overcorrection.py | 22 +- tests/test_public_api.py | 4 +- tests/test_sspsir.py | 171 +++++++------- 10 files changed, 389 insertions(+), 350 deletions(-) create mode 100644 docs/changes/devel/90.feature.rst diff --git a/docs/api.rst b/docs/api.rst index 5f17df8c..f8e85d06 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -128,8 +128,6 @@ SSP-SIR :nosignatures: mne_denoise.sspsir.SSPSIR - mne_denoise.sspsir.compute_sspsir - mne_denoise.sspsir.compute_sir Overcorrection metrics ---------------------- diff --git a/docs/changes/devel/90.feature.rst b/docs/changes/devel/90.feature.rst new file mode 100644 index 00000000..40998c7e --- /dev/null +++ b/docs/changes/devel/90.feature.rst @@ -0,0 +1,2 @@ +SSP-SIR now delegates Forward-based source-informed projection reconstruction +to MNE-Python. diff --git a/examples/sspsir/plot_01_sspsir_basics.py b/examples/sspsir/plot_01_sspsir_basics.py index 2baad251..ddf87dbb 100644 --- a/examples/sspsir/plot_01_sspsir_basics.py +++ b/examples/sspsir/plot_01_sspsir_basics.py @@ -12,7 +12,8 @@ space whose channels no longer correspond to electrodes. SSP-SIR repairs both problems: it reconstructs the projected data through a forward model, which restores what the head model says must have been there and returns the signal -to interpretable sensor space. +to interpretable sensor space. mne-denoise estimates the artifact subspace and +crossfade, while MNE-Python performs the Forward-based reconstruction. This example shows the two design choices that matter most in practice: how many artifact components to remove, and the fact that the projection is @@ -34,9 +35,10 @@ # %% # A montage and a forward model # ----------------------------- -# As with SOUND, the lead field is built from the montage when no ``forward`` -# is supplied. We build it explicitly here too, so the simulated brain -# response is something the head model can actually account for. +# As with SOUND, a spherical Forward is built from the montage when no +# ``forward`` is supplied. We build its gain explicitly here too, so the +# simulated brain response is something the head model can actually account +# for; SSP-SIR delegates its reconstruction to MNE-Python. montage = mne.channels.make_standard_montage("standard_1020") ch_names = [ch for ch in montage.ch_names if ch not in ("A1", "A2")][:32] diff --git a/mne_denoise/_leadfield.py b/mne_denoise/_leadfield.py index 995cfd58..0dad392a 100644 --- a/mne_denoise/_leadfield.py +++ b/mne_denoise/_leadfield.py @@ -1,31 +1,25 @@ -"""Lead-field construction shared by the forward-model denoisers. - -SOUND and SSP-SIR are both forward-model methods: they need a lead-field matrix -``L`` whose columns are the scalp topographies that cortical current sources -produce, and they use it as a prior for what a plausible brain signal looks -like. Whatever a sensor records that ``L`` cannot explain is treated as noise -or artifact, so the lead field is what separates "brain" from "everything -else" in both algorithms. - -An individualised forward model computed from the participant's anatomy is -always the better input. When none is available, both methods fall back to a -three-layer spherical head model derived from the electrode montage alone -(Mutanen et al., 2016, 2018). That is workable because neither algorithm uses -``L`` directly — only ``L @ L.T``, the lead-field covariance describing the -typical cross-correlations between channels — so a head model needs to capture -those correlations, not the anatomy that produced them. - -This module centralises the construction so the two estimators resolve their -lead field identically: +"""Forward and lead-field construction shared by the forward-model denoisers. + +SOUND consumes an average-referenced lead-field matrix ``L`` whose columns are +the scalp topographies that cortical current sources produce. SSP-SIR keeps the +MNE ``Forward`` used to produce those topographies and delegates its +projection-reconstruction mapping to MNE-Python. Both methods share the same +spherical ``Forward`` construction when an anatomy-specific model is +unavailable. + +This module centralises the construction so SOUND and SSP-SIR use the same +forward-model geometry: - :class:`SphericalHeadModel` — the shell radii and conductivities of the fallback head model, with the published values as :data:`REFERENCE_HEAD`. - :func:`fibonacci_sphere` — deterministic, quasi-uniform directions on a sphere, used to place the source dipoles. +- :func:`resolve_forward` — align a user-supplied ``Forward`` or build the + spherical fallback. - :func:`make_spherical_leadfield` — build the fallback lead field from a montage. - :func:`resolve_leadfield` — choose between a user-supplied forward model and - that fallback; this is the entry point the estimators call. + that fallback as an ndarray for SOUND. Every lead field returned here is average referenced, the reference both algorithms operate in. @@ -63,6 +57,7 @@ "SphericalHeadModel", "fibonacci_sphere", "make_spherical_leadfield", + "resolve_forward", "resolve_leadfield", ] @@ -173,31 +168,6 @@ def _forward_gain(forward: mne.Forward) -> np.ndarray: ) -def _leadfield_from_forward(forward: mne.Forward, info: mne.Info) -> np.ndarray: - """Extract an average-referenced lead field from a user forward solution. - - The forward's rows are reordered to match ``info``'s channel order, so a - forward computed elsewhere (with its own channel ordering) lines up with - the data being cleaned. - """ - gain = _forward_gain(forward) - row_names = list(forward["sol"]["row_names"]) - if len(row_names) != gain.shape[0]: - raise ValueError( - "The supplied forward has a different number of row names and " - "gain-matrix rows." - ) - wanted = list(info["ch_names"]) - missing = [ch for ch in wanted if ch not in row_names] - if missing: - raise ValueError( - "The supplied forward model is missing channels present in the data: " - f"{missing[:5]}{'...' if len(missing) > 5 else ''}." - ) - idx = [row_names.index(ch) for ch in wanted] - return _average_reference(gain[idx]) - - def fibonacci_sphere(n_points: int) -> np.ndarray: """Generate unit vectors quasi-uniformly covering the sphere. @@ -237,6 +207,73 @@ def fibonacci_sphere(n_points: int) -> np.ndarray: ) +def _make_spherical_forward( + info: mne.Info, + *, + n_dipoles: int = 5000, + head_model: SphericalHeadModel = REFERENCE_HEAD, + verbose: bool = False, +) -> mne.Forward: + """Build the spherical fallback ``Forward`` from an EEG montage.""" + if ( + isinstance(n_dipoles, (bool, np.bool_)) + or not isinstance(n_dipoles, Integral) + or n_dipoles < 1 + ): + raise ValueError(f"n_dipoles must be a positive integer, got {n_dipoles!r}.") + _mne.require_mne("automatic spherical lead-field construction") + eeg_picks = _mne.mne.pick_types(info, meg=False, eeg=True, exclude=()) + if len(eeg_picks) != len(info["ch_names"]): + raise ValueError( + "Automatic spherical lead-field construction supports EEG channels " + "only; provide an explicit forward model for MEG or mixed channel types." + ) + + with warnings.catch_warnings(): + # The best-fit sphere centre can sit >20 mm from the head-frame origin + # for partial or idealised montages; harmless for spanning the + # topography subspace that SOUND and SSP-SIR rely on. + warnings.filterwarnings("ignore", message=".*from head frame origin.*") + sphere = _mne.mne.make_sphere_model( + r0="auto", + head_radius="auto", + info=info, + relative_radii=head_model.relative_radii, + sigmas=head_model.conductivities, + verbose=verbose, + ) + # ``sphere["layers"][-1]["rad"]`` is the fitted scalp radius in + # metres; the source shell tracks it so the geometry stays + # proportional. + head_radius = float(sphere["layers"][-1]["rad"]) + directions = fibonacci_sphere(n_dipoles) + positions = ( + directions * (head_model.dipole_relative_radius * head_radius) + + sphere["r0"] + ) + src = _mne.mne.setup_volume_source_space( + pos={"rr": positions, "nn": directions}, sphere_units="m", verbose=verbose + ) + forward = _mne.mne.make_forward_solution( + info, + trans=None, + src=src, + bem=sphere, + eeg=True, + meg=False, + verbose=verbose, + ) + forward = _mne.mne.convert_forward_solution( + forward, force_fixed=True, use_cps=False, verbose=verbose + ) + # MNE's forward solver stores gains as float32. Keep the canonical + # Forward in double precision so its public gain view and MNE's Gram-based + # reconstruction use the same values. + forward["sol"]["data"] = np.asarray(forward["sol"]["data"], dtype=float) + forward["_orig_sol"] = np.asarray(forward["_orig_sol"], dtype=float) + return forward + + def make_spherical_leadfield( info: mne.Info, *, @@ -275,12 +312,12 @@ def make_spherical_leadfield( Notes ----- - Only ``leadfield @ leadfield.T`` enters SOUND and SSP-SIR, where it acts as - the source-covariance prior: it sets the minimum-norm weighting in SOUND - and the truncation scale of the source-informed reconstruction in SSP-SIR. - Its *spectrum* therefore matters, not merely its span, which is why the + SOUND uses ``leadfield @ leadfield.T`` as its source-covariance prior. SSP-SIR + retains the Forward and passes its gain, projections, whitening, and + source-informed reconstruction to MNE-Python. The *spectrum* of the gain + still matters for both methods, not merely its span, which is why the published shell-of-radial-dipoles geometry is followed rather than a volume - grid. A volume grid with free orientations spans a comparable subspace + grid. A volume grid with free orientations spans a comparable subspace (mean principal-angle cosine ~0.94 over the leading topographies for a 32-channel montage) but has a visibly faster-decaying spectrum, changing the effective regularisation. @@ -291,51 +328,67 @@ def make_spherical_leadfield( from such a draw — inside that sampling spread, while being exactly repeatable. """ - if ( - isinstance(n_dipoles, (bool, np.bool_)) - or not isinstance(n_dipoles, Integral) - or n_dipoles < 1 - ): - raise ValueError(f"n_dipoles must be a positive integer, got {n_dipoles!r}.") - _mne.require_mne("automatic spherical lead-field construction") - eeg_picks = _mne.mne.pick_types(info, meg=False, eeg=True, exclude=()) - if len(eeg_picks) != len(info["ch_names"]): - raise ValueError( - "Automatic spherical lead-field construction supports EEG channels " - "only; provide an explicit forward model for MEG or mixed channel types." - ) + forward = _make_spherical_forward( + info, n_dipoles=n_dipoles, head_model=head_model, verbose=verbose + ) + return _average_reference(_forward_gain(forward)) - with warnings.catch_warnings(): - # The best-fit sphere centre can sit >20 mm from the head-frame origin - # for partial or idealised montages; harmless for spanning the - # topography subspace that SOUND and SSP-SIR rely on. - warnings.filterwarnings("ignore", message=".*from head frame origin.*") - sphere = _mne.mne.make_sphere_model( - r0="auto", - head_radius="auto", - info=info, - relative_radii=head_model.relative_radii, - sigmas=head_model.conductivities, - verbose=verbose, - ) - # ``sphere["layers"][-1]["rad"]`` is the fitted scalp radius in metres; - # the source shell tracks it so the geometry stays proportional. - head_radius = float(sphere["layers"][-1]["rad"]) - directions = fibonacci_sphere(n_dipoles) - positions = ( - directions * (head_model.dipole_relative_radius * head_radius) - + sphere["r0"] - ) - src = _mne.mne.setup_volume_source_space( - pos={"rr": positions, "nn": directions}, sphere_units="m", verbose=verbose - ) - fwd = _mne.mne.make_forward_solution( - info, trans=None, src=src, bem=sphere, eeg=True, meg=False, verbose=verbose + +def resolve_forward( + *, + inst: mne.io.BaseRaw | mne.BaseEpochs | mne.Evoked | None, + ch_names: list[str] | None, + n_channels: int, + method: str, + forward: mne.Forward | None = None, + n_dipoles: int = 5000, + head_model: SphericalHeadModel = REFERENCE_HEAD, +) -> mne.Forward: + """Resolve and align the ``Forward`` used by a forward-model denoiser. + + MNE-Python performs the channel-name alignment for explicit forwards. For + array input, the forward's own row order is the channel contract because + the array carries no names. + """ + _mne.require_mne(f"{method} Forward resolution") + if inst is not None: + info = inst.copy().pick(ch_names).info + if forward is None: + return _make_spherical_forward( + info, n_dipoles=n_dipoles, head_model=head_model + ) + _forward_gain(forward) + resolved = _mne.mne.pick_channels_forward( + forward, include=info["ch_names"], ordered=True ) - fwd = _mne.mne.convert_forward_solution( - fwd, force_fixed=True, use_cps=False, verbose=verbose + resolved["sol"]["data"] = np.asarray(resolved["sol"]["data"], dtype=float) + resolved["_orig_sol"] = np.asarray(resolved["_orig_sol"], dtype=float) + return resolved + + if forward is not None: + gain = _forward_gain(forward) + if gain.shape[0] != n_channels: + raise ValueError( + "For array input, the forward must have the same number of " + f"channels as the data ({gain.shape[0]} vs {n_channels})." + ) + row_names = list(forward["sol"]["row_names"]) + if len(row_names) != gain.shape[0]: + raise ValueError( + "The supplied forward has a different number of row names and " + "gain-matrix rows." + ) + resolved = _mne.mne.pick_channels_forward( + forward, include=row_names, ordered=True ) - return _average_reference(np.asarray(fwd["sol"]["data"], dtype=float)) + resolved["sol"]["data"] = np.asarray(resolved["sol"]["data"], dtype=float) + resolved["_orig_sol"] = np.asarray(resolved["_orig_sol"], dtype=float) + return resolved + + raise ValueError( + f"{method} needs channel positions: pass an MNE object with a montage, " + "or provide a `forward` for array input." + ) def resolve_leadfield( @@ -350,9 +403,9 @@ def resolve_leadfield( ) -> np.ndarray: """Resolve the lead field an estimator should use. - Both :class:`~mne_denoise.sound.SOUND` and - :class:`~mne_denoise.sspsir.SSPSIR` accept the same three kinds of input, - and this function is where that choice is made once for both: + SOUND's ndarray lead-field input can be resolved from the same forward + model used by SSP-SIR. This function handles the three supported input + cases: - **MNE object, no forward** — build the spherical fallback from its montage. @@ -407,23 +460,13 @@ def resolve_leadfield( -------- make_spherical_leadfield : The fallback this dispatches to. """ - if inst is not None: - _mne.require_mne("MNE lead-field resolution") - info = inst.copy().pick(ch_names).info - if forward is not None: - return _leadfield_from_forward(forward, info) - return make_spherical_leadfield( - info, n_dipoles=n_dipoles, head_model=head_model - ) - if forward is not None: - gain = _forward_gain(forward) - if gain.shape[0] != n_channels: - raise ValueError( - "For array input, the forward must have the same number of " - f"channels as the data ({gain.shape[0]} vs {n_channels})." - ) - return _average_reference(gain) - raise ValueError( - f"{method} needs channel positions: pass an MNE object with a montage, " - "or provide a `forward` for array input." + forward = resolve_forward( + inst=inst, + ch_names=ch_names, + n_channels=n_channels, + method=method, + forward=forward, + n_dipoles=n_dipoles, + head_model=head_model, ) + return _average_reference(_forward_gain(forward)) diff --git a/mne_denoise/sspsir.py b/mne_denoise/sspsir.py index a7be5abb..231e1474 100644 --- a/mne_denoise/sspsir.py +++ b/mne_denoise/sspsir.py @@ -31,7 +31,6 @@ from __future__ import annotations -import warnings from numbers import Integral, Real import numpy as np @@ -43,7 +42,7 @@ from . import _mne from ._data import extract_data_from_mne, reconstruct_mne_object -from ._leadfield import _validate_leadfield, resolve_leadfield +from ._leadfield import _average_reference, _forward_gain, resolve_forward from ._logging import logger, verbose from ._validation import ( check_channel_layout, @@ -52,35 +51,12 @@ check_positive_real, ) -__all__ = ["SSPSIR", "compute_sir", "compute_sspsir"] +__all__ = ["SSPSIR"] #: 10-90% transition width (s) of the crossfade around a user artifact window. _SMOOTH_LENGTH = 0.010 -def _truncated_svd( - matrix: np.ndarray, M: int, what: str -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Return the leading singular triplets, limited to numerical rank.""" - if isinstance(M, (bool, np.bool_)) or not isinstance(M, Integral) or M < 1: - raise ValueError(f"M must be a positive integer, got {M!r}.") - u, s, vt = np.linalg.svd(matrix, full_matrices=False) - requested_rank = int(M) - tol = s[0] * max(matrix.shape) * np.finfo(s.dtype).eps if s.size else 0.0 - numerical_rank = int(np.count_nonzero(s > tol)) - if numerical_rank == 0: - raise ValueError(f"Cannot reconstruct from {what}: its numerical rank is zero.") - if numerical_rank < requested_rank: - warnings.warn( - f"M={M} exceeds the numerical rank ({numerical_rank}) of {what}; " - f"using M={numerical_rank} instead.", - RuntimeWarning, - stacklevel=3, - ) - rank = min(requested_rank, numerical_rank) - return u[:, :rank], s[:rank], vt[:rank] - - def _artifact_subspace( svd_input: np.ndarray, n_components ) -> tuple[np.ndarray, int, np.ndarray]: @@ -136,85 +112,6 @@ def _artifact_subspace( return u[:, :n_pc], n_pc, s -def compute_sspsir( - leadfield: np.ndarray, artifact_topographies: np.ndarray, M: int -) -> np.ndarray: - """Build the artifact-suppressing SSP-SIR operator (``cleaned = C @ data``). - - This is the projected branch of SSP-SIR: project the artifact subspace out, - then reconstruct through the forward model. In the full method it is - crossfaded against :func:`compute_sir`; see :class:`SSPSIR`. - - Parameters - ---------- - leadfield : ndarray, shape (n_channels, n_sources) - Average-referenced lead field. - artifact_topographies : ndarray, shape (n_channels, n_components) - Orthonormal artifact subspace (left singular vectors). - M : int - Truncation dimension of the source-informed reconstruction. - - Returns - ------- - operator : ndarray, shape (n_channels, n_channels) - The artifact-suppressing SSP-SIR operator. - """ - leadfield = _validate_leadfield(leadfield) - artifact_topographies = np.asarray(artifact_topographies, dtype=float) - if artifact_topographies.ndim != 2: - raise ValueError( - "artifact_topographies must be 2D, got shape " - f"{artifact_topographies.shape}." - ) - if not np.isfinite(artifact_topographies).all(): - raise ValueError("artifact_topographies must contain only finite values.") - n_channels = leadfield.shape[0] - if artifact_topographies.shape[0] != n_channels: - raise ValueError( - f"Lead field has {n_channels} channels but the artifact " - f"topographies have {artifact_topographies.shape[0]}." - ) - n_artifact = artifact_topographies.shape[1] - if not 1 <= n_artifact < n_channels: - raise ValueError( - "artifact_topographies must contain between 1 and n_channels - 1 " - "components." - ) - gram = artifact_topographies.T @ artifact_topographies - if not np.allclose(gram, np.eye(n_artifact), rtol=1e-7, atol=1e-9): - raise ValueError("artifact_topographies must have orthonormal columns.") - proj = np.eye(n_channels) - artifact_topographies @ artifact_topographies.T - pl = proj @ leadfield - u, s, vt = _truncated_svd(pl, M, "the projected lead field") - return leadfield @ (vt.T / s) @ u.T @ proj - - -def compute_sir(leadfield: np.ndarray, M: int) -> np.ndarray: - """Build the source-informed reconstruction operator without projection. - - A rank-``M`` reconstruction of the data through the forward model, with no - artifact subspace removed. It is - *not* the identity -- it restricts the data to the ``M`` leading - lead-field topographies -- so the crossfade in :class:`SSPSIR` applies the - same rank truncation inside and outside the artifact window. - - Parameters - ---------- - leadfield : ndarray, shape (n_channels, n_sources) - Average-referenced lead field. - M : int - Truncation dimension of the source-informed reconstruction. - - Returns - ------- - operator : ndarray, shape (n_channels, n_channels) - The unprojected source-informed reconstruction operator. - """ - leadfield = _validate_leadfield(leadfield) - u, _, _ = _truncated_svd(leadfield, M, "the lead field") - return u @ u.T - - def _as_mne_projections(topographies: np.ndarray, ch_names: list[str]) -> list: """Wrap artifact topographies as :class:`mne.Projection` objects. @@ -243,21 +140,36 @@ def _as_mne_projections(topographies: np.ndarray, ch_names: list[str]) -> list: ] +def _reconstruction_operator(info, forward, projs, rank: int) -> np.ndarray: + """Build a channel-space reconstruction operator through MNE-Python.""" + _mne.require_mne("SSP-SIR projection reconstruction") + n_channels = len(info["ch_names"]) + identity = _mne.mne.EvokedArray( + np.eye(n_channels), info.copy(), tmin=0.0, verbose=False + ) + with identity.info._unlock(): + identity.info["projs"] = [] + identity.add_proj(projs, verbose=False) + return identity.reconstruct_proj(forward=forward, rank=rank).data + + class SSPSIR(BaseEstimator, TransformerMixin): """Suppress TMS-evoked muscle artifacts from EEG (SSP-SIR). SSP-SIR projects out the high-variance muscle-artifact subspace and then reconstructs the brain signal lost to that projection through a forward - model. With no individualised forward model, a three-layer spherical lead - field is built from the montage. Conceptually this is signal-space - projection followed by the source-informed reconstruction that MNE-Python - also exposes as ``proj="reconstruct"``. + model. With no individualised forward model, a three-layer spherical + ``Forward`` is built from the montage. Conceptually this is signal-space + projection followed by the source-informed reconstruction provided by + MNE-Python's ``Evoked.reconstruct_proj(forward=..., rank=...)``. + mne-denoise estimates the artifact subspace and temporal crossfade; + MNE-Python performs the Forward-based projection reconstruction. As in the published method, the artifact-suppressed reconstruction is - crossfaded in time against a reconstruction of the *unprojected* data - (:func:`compute_sir`), so that the artifact subspace is removed - only where the muscle artifact lives. Without this crossfade the projection - also removes brain signal from the baseline and from late TEP components. + crossfaded in time against an unprojected, rank-limited source-informed + reconstruction, so that the artifact subspace is removed only where the + muscle artifact lives. Without this crossfade the projection also removes + brain signal from the baseline and from late TEP components. Restricting suppression to the artifact window is also what the method's authors do in practice (Mutanen et al., 2024) [2]_. @@ -272,7 +184,8 @@ class SSPSIR(BaseEstimator, TransformerMixin): cumulative high-frequency variance fraction to cover (float in (0, 1)). See Notes for the exact criterion. forward : mne.Forward | None - Optional forward solution; if None a spherical lead field is built. + Optional forward solution; if None a spherical ``Forward`` is built + from the fitted object's montage. art_window : tuple of float | None ``(tmin, tmax)`` in seconds delimiting the muscle artifact used to estimate the artifact subspace. If None, the subspace is estimated @@ -294,15 +207,19 @@ class SSPSIR(BaseEstimator, TransformerMixin): sfreq : float | None Sampling frequency, required only for plain-array input. n_dipoles : int - Number of dipoles for the spherical lead field. + Number of dipoles for the spherical fallback ``Forward``. verbose : bool | str | int | None, default=None MNE-style logging level. The fitted SSP-SIR summary is emitted at - INFO; numerical reconstruction helpers remain silent. + INFO; projection reconstruction is performed once during fitting. Attributes ---------- + forward_ : mne.Forward + Forward solution used for projection reconstruction, aligned to the + fitted channel order. leadfield_ : ndarray - The average-referenced lead field used. + Average-referenced gain matrix from ``forward_`` in fitted channel + order, retained for diagnostics and overcorrection metrics. artifact_topographies_ : ndarray, shape (n_channels, n_components) The removed artifact subspace. operator_ : ndarray, shape (n_channels, n_channels) @@ -315,17 +232,16 @@ class SSPSIR(BaseEstimator, TransformerMixin): n_components_ : int Number of artifact components removed. M_ : int - Effective source-informed reconstruction rank used. This can be lower - than the requested ``M`` when the projected lead field has lower - numerical rank. + Effective source-informed reconstruction rank of ``operator_`` after + MNE-Python's exact-rank reconstruction. singular_values_ : ndarray Singular values of the high-frequency data the subspace was estimated from. Inspect these to choose ``n_components`` by the spectrum elbow, which is what Mutanen et al. (2016) actually recommend [1]_. projs_ : list of mne.Projection ``artifact_topographies_`` wrapped as MNE projections, so the removed - directions can be plotted with ``mne.viz.plot_projs_topomap``. Empty - when fitted on a plain array, which carries no channel names. + directions can be plotted with ``mne.viz.plot_projs_topomap``. For + array input, channel names are taken from the fitted Forward. Notes ----- @@ -519,7 +435,7 @@ def fit( self.singular_values_, ) = _artifact_subspace(svd_input, self.n_components) - self.leadfield_ = resolve_leadfield( + self.forward_ = resolve_forward( inst=orig_inst, ch_names=ch_names, n_channels=n_channels, @@ -527,12 +443,37 @@ def fit( forward=self.forward, n_dipoles=self.n_dipoles, ) + fitted_ch_names = ( + list(ch_names) + if ch_names is not None + else list(self.forward_["sol"]["row_names"]) + ) + self.leadfield_ = _average_reference(_forward_gain(self.forward_)) + if orig_inst is not None: + fit_info = orig_inst.copy().pick(fitted_ch_names).info + else: + fit_info = self.forward_["info"].copy() + # Forward metadata has channel geometry but does not necessarily + # carry the sampling frequency needed by EvokedArray. + with fit_info._unlock(): + fit_info["sfreq"] = sfreq + self.projs_ = _as_mne_projections(self.artifact_topographies_, fitted_ch_names) + from mne.proj import make_eeg_average_ref_proj + + average_ref = make_eeg_average_ref_proj(fit_info, verbose=False) data_rank = int(np.linalg.matrix_rank(evoked)) - M = self.M if self.M is not None else max(1, data_rank - self.n_components_) - self.operator_ = compute_sspsir(self.leadfield_, self.artifact_topographies_, M) + M = int(self.M) if self.M is not None else data_rank - self.n_components_ + self.operator_ = _reconstruction_operator( + fit_info, + self.forward_, + [average_ref, *self.projs_], + M, + ) self.M_ = int(np.linalg.matrix_rank(self.operator_)) - self.operator_orig_ = compute_sir(self.leadfield_, self.M_) + self.operator_orig_ = _reconstruction_operator( + fit_info, self.forward_, [average_ref], self.M_ + ) self.kernel_ = ( np.ones(evoked.shape[1]) if self.blend == "constant" @@ -540,12 +481,7 @@ def fit( ) self.sfreq_ = sfreq self.times_ = times.copy() - self._mne_ch_names_ = ch_names - self.projs_ = ( - _as_mne_projections(self.artifact_topographies_, ch_names) - if ch_names is not None - else [] - ) + self._mne_ch_names_ = fitted_ch_names logger.info( "SSP-SIR: channels=%d, removed %d artifact component(s), " "SIR truncation M=%d (data rank %d), blend=%s", diff --git a/tests/test_leadfield.py b/tests/test_leadfield.py index 654d2d95..4e497a7c 100644 --- a/tests/test_leadfield.py +++ b/tests/test_leadfield.py @@ -9,9 +9,11 @@ from mne_denoise._leadfield import ( REFERENCE_HEAD, SphericalHeadModel, + _make_spherical_forward, _validate_leadfield, fibonacci_sphere, make_spherical_leadfield, + resolve_forward, resolve_leadfield, ) @@ -65,6 +67,21 @@ def test_spherical_leadfield_is_deterministic(eeg_info): np.testing.assert_array_equal(a, b) +def test_spherical_forward_is_fixed_deterministic_and_has_requested_sources( + eeg_info, +): + """The shared fallback returns the fixed radial Forward itself.""" + first = _make_spherical_forward(eeg_info, n_dipoles=40) + second = _make_spherical_forward(eeg_info, n_dipoles=40) + + assert isinstance(first, mne.Forward) + assert mne.forward.is_fixed_orient(first) + assert first["sol"]["data"].shape == (24, 40) + assert len(first["src"][0]["vertno"]) == 40 + np.testing.assert_array_equal(first["src"][0]["rr"], second["src"][0]["rr"]) + np.testing.assert_array_equal(first["src"][0]["nn"], second["src"][0]["nn"]) + + def test_spherical_leadfield_matches_reference_geometry(): """Spectrum of L @ L.T matches the reference random-shell construction. @@ -211,6 +228,38 @@ def test_resolve_leadfield_from_forward(eeg_info, forward): assert np.allclose(leadfield.mean(axis=0), 0.0, atol=1e-9) +def test_resolve_forward_aligns_explicit_forward(eeg_info, forward): + """The canonical Forward is aligned by MNE's channel picker.""" + names = list(eeg_info["ch_names"])[::-1] + flipped_info = mne.create_info(names, 1000.0, "eeg") + flipped_info.set_montage("standard_1020") + + resolved = resolve_forward( + inst=_raw(flipped_info), + ch_names=names, + n_channels=24, + method="SSP-SIR", + forward=forward, + ) + assert resolved["sol"]["row_names"] == names + np.testing.assert_allclose( + resolved["sol"]["data"], forward["sol"]["data"][::-1], atol=1e-12 + ) + + +def test_resolve_forward_array_uses_forward_order(forward): + """Array input uses the Forward metadata as its channel contract.""" + resolved = resolve_forward( + inst=None, + ch_names=None, + n_channels=24, + method="SSP-SIR", + forward=forward, + ) + assert resolved["sol"]["row_names"] == forward["sol"]["row_names"] + assert resolved["info"]["ch_names"] == resolved["sol"]["row_names"] + + def test_resolve_leadfield_aligns_forward_channels(eeg_info, forward): """Reordering the data's channels reorders the forward's rows to match.""" names = list(eeg_info["ch_names"]) @@ -239,8 +288,8 @@ def test_resolve_leadfield_forward_missing_channels_raises(forward): ch = mne.channels.make_standard_montage("standard_1020").ch_names[:26] info = mne.create_info(ch, 1000.0, "eeg") info.set_montage("standard_1020") - with pytest.raises(ValueError, match="missing channels"): - resolve_leadfield( + with pytest.raises(ValueError, match="Missing channels"): + resolve_forward( inst=_raw(info), ch_names=info["ch_names"], n_channels=26, diff --git a/tests/test_mne.py b/tests/test_mne.py index e00dda0f..f5716fbe 100644 --- a/tests/test_mne.py +++ b/tests/test_mne.py @@ -51,7 +51,6 @@ def blocked_import(name, *args, **kwargs): from mne_denoise.sns import compute_sns from mne_denoise.sound import compute_sound, compute_sound_ref_best from mne_denoise.spectrum_interpolation import interpolate_spectrum - from mne_denoise.sspsir import compute_sir, compute_sspsir from mne_denoise.ssa import compute_basic_ssa rng = np.random.default_rng(0) @@ -66,9 +65,6 @@ def blocked_import(name, *args, **kwargs): compute_sound(data, leadfield, n_iter=1, random_state=0) compute_sound_ref_best(data, leadfield, n_iter=1, random_state=0) - compute_sir(leadfield, 2) - artifact_basis, _ = np.linalg.qr(rng.standard_normal((4, 1))) - compute_sspsir(leadfield, artifact_basis, 1) compute_dss(np.eye(4), np.diag([4.0, 3.0, 2.0, 1.0]), n_components=2) iterative_dss(data, lambda source: source**3, 1, max_iter=2, random_state=0) compute_sns(data, n_neighbors=2) diff --git a/tests/test_overcorrection.py b/tests/test_overcorrection.py index cd56b918..831b6570 100644 --- a/tests/test_overcorrection.py +++ b/tests/test_overcorrection.py @@ -118,15 +118,31 @@ def test_works_on_a_fitted_sspsir_operator(): import mne from mne_denoise._leadfield import make_spherical_leadfield - from mne_denoise.sspsir import compute_sspsir + from mne_denoise.sspsir import SSPSIR names = mne.channels.make_standard_montage("standard_1020").ch_names[:24] info = mne.create_info(names, 1000.0, "eeg") info.set_montage("standard_1020") lf = make_spherical_leadfield(info, n_dipoles=200) topos = np.linalg.svd(lf, full_matrices=False)[0][:, :2] - - m = quantify_overcorrection(compute_sspsir(lf, topos, M=15), lf) + times = np.arange(400) / 1000.0 + sources = np.stack( + [ + np.sin(2.0 * np.pi * 180.0 * times), + np.cos(2.0 * np.pi * 210.0 * times), + ] + ) + data = topos @ sources + evoked = mne.EvokedArray(data * 1e-6, info, tmin=0.0) + + ss = SSPSIR( + n_components=2, + M=15, + n_dipoles=200, + art_window=(0.0, 0.399), + blend="constant", + ).fit(evoked) + m = quantify_overcorrection(ss.operator_, ss.leadfield_) assert m["correlation"].shape == (200,) # Removing two dimensions should preserve most sources but not all. assert 0.0 < np.nanmean(m["goodness_of_fit"]) < 1.0 diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 248cc8f5..f8a985fd 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -49,7 +49,7 @@ def test_flattened_method_modules_expose_canonical_api(): SpectrumInterpolation, interpolate_spectrum, ) - from mne_denoise.sspsir import SSPSIR, compute_sir, compute_sspsir + from mne_denoise.sspsir import SSPSIR assert all( callable(item) @@ -67,8 +67,6 @@ def test_flattened_method_modules_expose_canonical_api(): SpectrumInterpolation, interpolate_spectrum, SSPSIR, - compute_sir, - compute_sspsir, ) ) diff --git a/tests/test_sspsir.py b/tests/test_sspsir.py index b8f1f430..2c524ec6 100644 --- a/tests/test_sspsir.py +++ b/tests/test_sspsir.py @@ -7,7 +7,7 @@ import pytest from sklearn.exceptions import NotFittedError -from mne_denoise.sspsir import SSPSIR, _artifact_subspace, compute_sir, compute_sspsir +from mne_denoise.sspsir import SSPSIR, _artifact_subspace @pytest.fixture(scope="module") @@ -246,6 +246,21 @@ def test_artifact_subspace_rejects_invalid_input(data, match): _artifact_subspace(data, 1) +def _direct_sspsir_operator(leadfield, artifact_topographies, rank): + """Independent direct SSP-SIR formula used as a parity oracle.""" + projector = ( + np.eye(leadfield.shape[0]) - artifact_topographies @ artifact_topographies.T + ) + u, s, vh = np.linalg.svd(projector @ leadfield, full_matrices=False) + return leadfield @ (vh[:rank].T / s[:rank]) @ u[:, :rank].T @ projector + + +def _direct_sir_operator(leadfield, rank): + """Independent direct rank-limited SIR formula used as a parity oracle.""" + u, _, _ = np.linalg.svd(leadfield, full_matrices=False) + return u[:, :rank] @ u[:, :rank].T + + def test_sspsir_manual_window(tms_epochs): epochs = tms_epochs[0] ss = SSPSIR(n_components=2, art_window=(0.005, 0.050)).fit(epochs) @@ -326,89 +341,9 @@ def test_sspsir_not_fitted_raises(): SSPSIR(n_components=2).transform(np.zeros((24, 100))) -def test_compute_sspsir_shape(eeg_info): - from mne_denoise._leadfield import make_spherical_leadfield - - leadfield = make_spherical_leadfield(eeg_info, n_dipoles=300) - u = np.linalg.svd(np.random.default_rng(0).standard_normal((24, 24)))[0][:, :2] - operator = compute_sspsir(leadfield, u, M=20) - assert operator.shape == (24, 24) - - -def test_compute_sspsir_channel_mismatch_raises(): - leadfield = np.random.default_rng(0).standard_normal((24, 100)) - bad = np.linalg.svd(np.random.default_rng(1).standard_normal((30, 30)))[0][:, :2] - with pytest.raises(ValueError, match="artifact"): - compute_sspsir(leadfield, bad, M=10) - - -def test_compute_sir_uses_shared_leadfield_validator(monkeypatch): - import mne_denoise.sspsir as sspsir_core - - leadfield = np.random.default_rng(12).standard_normal((5, 8)) - seen = [] - original = sspsir_core._validate_leadfield - - def wrapped(value, **kwargs): - seen.append(value) - return original(value, **kwargs) - - monkeypatch.setattr(sspsir_core, "_validate_leadfield", wrapped) - compute_sir(leadfield, M=3) - - assert len(seen) == 1 - assert seen[0] is leadfield - - -def test_compute_sspsir_rejects_nonorthonormal_subspace(): - leadfield = np.random.default_rng(11).standard_normal((5, 10)) - with pytest.raises(ValueError, match="orthonormal"): - compute_sspsir(leadfield, np.ones((5, 2)), M=3) - - -@pytest.mark.parametrize( - ("leadfield", "topographies", "M", "message"), - [ - (np.ones((4, 3)), np.ones(4), 2, "artifact_topographies must be 2D"), - (np.ones((4, 3)), np.full((4, 1), np.nan), 2, "finite"), - (np.ones((4, 3)), np.empty((4, 0)), 2, "between 1"), - (np.ones((4, 3)), np.eye(4)[:, :1], 0, "positive integer"), - ], -) -def test_compute_sspsir_input_contracts(leadfield, topographies, M, message): - with pytest.raises(ValueError, match=message): - compute_sspsir(leadfield, topographies, M) - - -def test_compute_sir_is_rank_m_not_identity(): - """orig_data_SIR restricts to the M leading topographies; it is not I.""" - leadfield = np.random.default_rng(2).standard_normal((24, 100)) - leadfield -= leadfield.mean(axis=0, keepdims=True) - operator = compute_sir(leadfield, M=15) - assert np.linalg.matrix_rank(operator) == 15 - assert not np.allclose(operator, np.eye(24)) - - -def test_compute_sir_rejects_invalid_m_or_zero_rank(): - with pytest.raises(ValueError, match="positive integer"): - compute_sir(np.eye(3), M=1.5) - with pytest.raises(ValueError, match="rank is zero"): - compute_sir(np.zeros((3, 4)), M=1) - - -def test_sspsir_warns_when_m_exceeds_rank(): - leadfield = np.random.default_rng(3).standard_normal((24, 100)) - leadfield -= leadfield.mean(axis=0, keepdims=True) # rank 23 - u = np.linalg.svd(np.random.default_rng(4).standard_normal((24, 24)))[0][:, :2] - with pytest.warns(RuntimeWarning, match="exceeds the numerical rank"): - operator = compute_sspsir(leadfield, u, M=24) - assert np.isfinite(operator).all() - - def test_sspsir_records_effective_m(tms_epochs): - with pytest.warns(RuntimeWarning, match="exceeds the numerical rank"): - ss = SSPSIR(n_components=2, M=100).fit(tms_epochs[0]) - assert ss.M_ < 100 + ss = SSPSIR(n_components=2, M=10).fit(tms_epochs[0]) + assert ss.M_ == 10 assert ss.M_ == np.linalg.matrix_rank(ss.operator_) assert ss.M_ == np.linalg.matrix_rank(ss.operator_orig_) @@ -437,9 +372,72 @@ def test_sspsir_mne_object_with_forward(tms_epochs, forward): epochs = tms_epochs[0] ss = SSPSIR(n_components=2, forward=forward).fit(epochs) assert ss.leadfield_.shape == (24, forward["sol"]["data"].shape[1]) + evoked = epochs.get_data().mean(axis=0) + evoked -= evoked.mean(axis=0, keepdims=True) + assert ss.M_ == np.linalg.matrix_rank(evoked) - ss.n_components_ + assert ss.transform(epochs).get_data().shape == epochs.get_data().shape + + +def test_sspsir_spherical_forward_parity_and_transform(tms_epochs): + """The spherical fallback keeps its Forward and matches direct SSP-SIR.""" + epochs = tms_epochs[0] + ss = SSPSIR(n_components=2, M=10, n_dipoles=300, blend="constant").fit(epochs) + + assert isinstance(ss.forward_, mne.Forward) + assert ss.forward_["sol"]["row_names"] == epochs.ch_names + expected_leadfield = ss.forward_["sol"]["data"] - ss.forward_["sol"]["data"].mean( + axis=0, keepdims=True + ) + np.testing.assert_allclose(ss.leadfield_, expected_leadfield, atol=1e-12) + np.testing.assert_allclose( + ss.operator_, + _direct_sspsir_operator(ss.leadfield_, ss.artifact_topographies_, ss.M_), + rtol=1e-12, + atol=1e-12, + ) + np.testing.assert_allclose( + ss.operator_orig_, + _direct_sir_operator(ss.leadfield_, ss.M_), + rtol=1e-12, + atol=1e-12, + ) assert ss.transform(epochs).get_data().shape == epochs.get_data().shape +def test_sspsir_explicit_forward_parity_and_channel_alignment(tms_epochs, forward): + """MNE reconstruction is invariant to the supplied Forward row order.""" + epochs = tms_epochs[0] + reversed_forward = mne.pick_channels_forward( + forward, include=forward["sol"]["row_names"][::-1], ordered=True + ) + straight = SSPSIR(n_components=2, M=10, forward=forward, blend="constant").fit( + epochs + ) + reversed_ = SSPSIR( + n_components=2, M=10, forward=reversed_forward, blend="constant" + ).fit(epochs) + + for ss in (straight, reversed_): + assert ss.forward_["sol"]["row_names"] == epochs.ch_names + np.testing.assert_allclose( + ss.operator_, + _direct_sspsir_operator(ss.leadfield_, ss.artifact_topographies_, ss.M_), + rtol=1e-12, + atol=1e-12, + ) + np.testing.assert_allclose( + ss.operator_orig_, + _direct_sir_operator(ss.leadfield_, ss.M_), + rtol=1e-12, + atol=1e-12, + ) + np.testing.assert_allclose(reversed_.leadfield_, straight.leadfield_, atol=1e-12) + np.testing.assert_allclose(reversed_.operator_, straight.operator_, atol=1e-12) + np.testing.assert_allclose( + reversed_.transform(epochs).get_data(), straight.transform(epochs).get_data() + ) + + def test_sspsir_array_forward_channel_mismatch_raises(forward): arr = np.random.default_rng(9).standard_normal((30, 400)) with pytest.raises(ValueError, match="same number of"): @@ -471,11 +469,12 @@ def test_sspsir_exposes_mne_projections(tms_epochs): mne_mod.viz.plot_projs_topomap(ss.projs_, epochs.info, show=False) -def test_sspsir_projs_empty_for_array_input(forward): - """A plain array has no channel names, so no projections can be built.""" +def test_sspsir_projs_use_forward_names_for_array_input(forward): + """Array input exposes projections using the Forward's channel names.""" arr = np.random.default_rng(11).standard_normal((24, 400)) ss = SSPSIR(n_components=2, sfreq=1000.0, forward=forward).fit(arr) - assert ss.projs_ == [] + assert len(ss.projs_) == ss.n_components_ + assert all(p["data"]["col_names"] == forward["sol"]["row_names"] for p in ss.projs_) def test_sspsir_verbose_logs_fit_summary(tms_epochs, caplog):