From 1661af26e14b30c7caf9fa5adec1037e8b8231fc Mon Sep 17 00:00:00 2001 From: Sina Esmaeili Date: Thu, 2 Jul 2026 06:40:34 -0400 Subject: [PATCH 1/4] ENH: add multi-channel Wiener filter (MWF) native denoiser Add mne_denoise.mwf with MWF (BaseEstimator/TransformerMixin) plus compute_mwf, mwf_filter, and hf_power_mask. MWF is a generic reference-free spatial cleaner and the RELAX-pipeline core (Somers, Francart & Bertrand 2018): it recovers the clean signal via R_clean @ inv(R_artifact) @ X from artifact-free and artifact-present segment covariances, with artifact segments marked by broadband HF power (or a supplied mask). fit() learns the Wiener operator, transform() applies it (leakage-safe); accepts MNE Raw/Epochs (sfreq from info) or NumPy arrays. Documented as a general cleaner (can attenuate neural HF activity), not an artifact-specific method. Edge cases handled: cap the HF-mask smoothing window at the signal length (numpy convolve mode=same otherwise yields a mask longer than short recordings/epochs); guard single-channel input; coerce int 0/1 masks to bool. Includes unit tests (tests/test_mwf.py) and docs (docs/mwf.md, api/index). --- docs/api.rst | 11 ++ docs/index.rst | 1 + docs/mwf.md | 83 +++++++++ mne_denoise/__init__.py | 7 +- mne_denoise/mwf/__init__.py | 31 +++ mne_denoise/mwf/core.py | 363 ++++++++++++++++++++++++++++++++++++ tests/test_mwf.py | 218 ++++++++++++++++++++++ 7 files changed, 713 insertions(+), 1 deletion(-) create mode 100644 docs/mwf.md create mode 100644 mne_denoise/mwf/__init__.py create mode 100644 mne_denoise/mwf/core.py create mode 100644 tests/test_mwf.py diff --git a/docs/api.rst b/docs/api.rst index e00aeaf5..046b9812 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -63,6 +63,17 @@ iCanClean mne_denoise.icanclean.ICanClean mne_denoise.icanclean.compute_icanclean +MWF +--- +.. autosummary:: + :toctree: generated/ + :nosignatures: + + mne_denoise.mwf.MWF + mne_denoise.mwf.compute_mwf + mne_denoise.mwf.hf_power_mask + mne_denoise.mwf.mwf_filter + Denoisers --------- .. autosummary:: diff --git a/docs/index.rst b/docs/index.rst index 3efae80a..687c8bc4 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -13,6 +13,7 @@ to extract reproducible or rhythmic components while preserving data rank. getting-started dss asr + mwf auto_examples/index .. toctree:: diff --git a/docs/mwf.md b/docs/mwf.md new file mode 100644 index 00000000..743b8f19 --- /dev/null +++ b/docs/mwf.md @@ -0,0 +1,83 @@ +# Multi-channel Wiener Filter (MWF) + +## Overview + +The `mne_denoise.mwf` module implements the **multi-channel Wiener filter (MWF)** +(Somers, Francart & Bertrand 2018), a generic, reference-free spatial artifact +cleaner and the spatial-filter core of the **RELAX** pipeline (Bailey et al. +2023). + +With the recording split into artifact-present and artifact-free segments, the +clean signal is recovered by + +``` +X_clean = R_clean · R_artifact⁻¹ · X +``` + +where `R_artifact` is the covariance over artifact segments (signal + artifact) +and `R_clean` the covariance over artifact-free segments (signal only). During +clean segments the two covariances coincide, so the filter is ~identity (no +over-cleaning); during artifact segments it projects out the artifact subspace. +No reference channel is required — artifact segments are marked from broadband +high-frequency power (or a caller-supplied mask). + +> **Note.** MWF is a *general* cleaner, not an artifact-specific method. Because +> its clean/artifact split is driven by broadband HF power, it can attenuate +> genuine neural high-frequency activity when the two covariances are poorly +> separated. Validate preservation of the band of interest on your data. It is +> provided here as the well-cited RELAX-core building block. + +## Quick Start + +```python +import numpy as np +from mne_denoise.mwf import MWF + +data = np.random.randn(32, 20000) # 32 channels, 20000 samples + +# Leakage-safe estimator API: learn the operator on train, apply to eval. +est = MWF(sfreq=250.0) +est.fit(data) +cleaned = est.transform(data) +print(f"fit on {est.artifact_fraction_:.0%} artifact samples") +``` + +Supplying an explicit artifact mask (no HF detector, no `sfreq` needed): + +```python +cleaned = MWF().fit_transform(data, mask=my_artifact_mask) +``` + +With MNE-Python objects the sampling frequency is read from `info`: + +```python +raw_clean = MWF().fit_transform(raw) +``` + +## One-shot functional API + +```python +from mne_denoise.mwf import compute_mwf, hf_power_mask, mwf_filter + +cleaned, info = compute_mwf(data, sfreq=250.0) # (cleaned, {mask, artifact_fraction}) +mask = hf_power_mask(data, sfreq=250.0) # broadband HF artifact detector +cleaned = mwf_filter(data, mask) # raw Wiener filter given a mask +``` + +## Parameters + +| Parameter | Description | +| ---------- | --------------------------------------------------------------------- | +| `sfreq` | Sampling frequency (Hz). Optional for MNE input / when a mask is given. | +| `hf_hz` | High-pass cutoff (Hz) for the HF artifact detector. | +| `quantile` | HF-power quantile above which samples are flagged as artifact. | +| `reg` | Diagonal-loading factor for covariance invertibility. | + +## References + +1. Somers, B., Francart, T., & Bertrand, A. (2018). A generic EEG artifact + removal algorithm based on the multi-channel Wiener filter. _Journal of Neural + Engineering_, 15(3), 036007. https://doi.org/10.1088/1741-2552/aaac92 +2. Bailey, N. W., et al. (2023). RELAX: An automated pre-processing pipeline for + cleaning EEG data — Part 1. _Clinical Neurophysiology_, 149, 178-201. + https://doi.org/10.1016/j.clinph.2023.01.007 diff --git a/mne_denoise/__init__.py b/mne_denoise/__init__.py index d916072a..27d5fee7 100644 --- a/mne_denoise/__init__.py +++ b/mne_denoise/__init__.py @@ -30,9 +30,13 @@ spectrum_interpolation : Spectrum Interpolation Removes power-line noise and its harmonics by interpolating spectral amplitudes while preserving phase. + +mwf : Multi-channel Wiener Filtering + Semi-supervised covariance-based artifact suppression using explicit artifact + and clean training segments. """ -from . import asr, dss, icanclean, spectrum_interpolation, zapline +from . import asr, dss, icanclean, mwf, spectrum_interpolation, zapline __version__ = "0.0.1" @@ -41,5 +45,6 @@ "dss", "icanclean", "spectrum_interpolation", + "mwf", "zapline", ] diff --git a/mne_denoise/mwf/__init__.py b/mne_denoise/mwf/__init__.py new file mode 100644 index 00000000..d6b7034d --- /dev/null +++ b/mne_denoise/mwf/__init__.py @@ -0,0 +1,31 @@ +"""Multi-channel Wiener filter (MWF) artifact removal. + +This module contains: + +- ``hf_power_mask``: broadband high-frequency artifact-segment detector. +- ``mwf_filter`` / ``compute_mwf``: the array-based Wiener filter. +- ``MWF``: the scikit-learn estimator, compatible with MNE-Python objects or + NumPy arrays. + +The MWF (Somers, Francart & Bertrand 2018) is a generic, reference-free spatial +artifact cleaner and the spatial-filter core of the RELAX pipeline. It recovers +the clean signal via ``R_clean @ R_artifact^{-1} @ X`` from artifact-free and +artifact-present segment covariances. It is a general cleaner rather than an +artifact-specific method; validate preservation of the band of interest. + +References +---------- +.. [1] Somers, B., Francart, T., & Bertrand, A. (2018). A generic EEG artifact + removal algorithm based on the multi-channel Wiener filter. Journal of + Neural Engineering, 15(3), 036007. + https://doi.org/10.1088/1741-2552/aaac92 +""" + +from .core import MWF, compute_mwf, hf_power_mask, mwf_filter + +__all__ = [ + "MWF", + "compute_mwf", + "hf_power_mask", + "mwf_filter", +] diff --git a/mne_denoise/mwf/core.py b/mne_denoise/mwf/core.py new file mode 100644 index 00000000..dc9caea4 --- /dev/null +++ b/mne_denoise/mwf/core.py @@ -0,0 +1,363 @@ +"""Multi-channel Wiener filter (MWF) for EEG artifact removal. + +The multi-channel Wiener filter (Somers, Francart & Bertrand 2018) [1]_ is a +generic, reference-free artifact remover and the spatial-filter core of the +RELAX pipeline (Bailey et al. 2023) [2]_. With the recording split into +artifact-present and artifact-free segments, the clean signal is recovered by + + X_clean = R_clean @ R_artifact^{-1} @ X + +where ``R_artifact`` is the covariance over artifact segments (signal + artifact) +and ``R_clean`` the covariance over artifact-free segments (signal only). During +clean segments the two covariances coincide so the filter is ~identity (no +over-cleaning); during artifact segments it projects out the artifact subspace. +No reference channel is needed: artifact segments are marked by a +high-frequency-power threshold (a caller-supplied mask is honoured). + +.. note:: + + MWF is a *general* spatial cleaner rather than an artifact-specific method. + Because its clean/artifact split is driven by broadband high-frequency power, + it can attenuate genuine neural high-frequency activity when the two + covariances are poorly separated; validate preservation of the band of + interest on your data. In our M/EEG benchmark it was less selective than the + reference-free CCA/SSA methods for muscle removal — it is provided here as the + well-cited RELAX-core building block, not as a targeted muscle remover. + +This module contains: + +- ``hf_power_mask``: broadband high-frequency artifact-segment detector. +- ``mwf_filter`` / ``compute_mwf``: the array-based Wiener filter. +- ``MWF``: the scikit-learn estimator, compatible with MNE-Python objects or + NumPy arrays. + +References +---------- +.. [1] Somers, B., Francart, T., & Bertrand, A. (2018). A generic EEG artifact + removal algorithm based on the multi-channel Wiener filter. Journal of + Neural Engineering, 15(3), 036007. + https://doi.org/10.1088/1741-2552/aaac92 +.. [2] Bailey, N. W., et al. (2023). RELAX: An automated pre-processing pipeline + for cleaning EEG data - Part 1. Clinical Neurophysiology, 149, 178-201. + https://doi.org/10.1016/j.clinph.2023.01.007 +""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np +from sklearn.base import BaseEstimator, TransformerMixin +from sklearn.utils.validation import check_is_fitted + +from ..utils import extract_data_from_mne, reconstruct_mne_object + +logger = logging.getLogger(__name__) + + +def hf_power_mask( + X: np.ndarray, + sfreq: float, + hf_hz: float = 20.0, + quantile: float = 0.6, + smooth_s: float = 0.1, +) -> np.ndarray: + """Mark high-frequency-power time points as artifact (broadband muscle/transient). + + Parameters + ---------- + X : ndarray, shape (n_channels, n_times) + Multichannel signal. + sfreq : float + Sampling frequency in Hz. + hf_hz : float + High-pass cutoff (Hz) used to isolate broadband high-frequency power. + quantile : float + Samples whose smoothed HF power exceeds this quantile are flagged. + smooth_s : float + Moving-average smoothing window (seconds) applied to the HF envelope. + + Returns + ------- + mask : ndarray of bool, shape (n_times,) + ``True`` where the sample is flagged as artifact. + """ + from scipy.signal import butter, filtfilt + + ny = 0.5 * float(sfreq) + b, a = butter(4, min(float(hf_hz) / ny, 0.99), btype="high") + hf = filtfilt(b, a, np.asarray(X, dtype=float), axis=-1) + env = (hf**2).mean(axis=0) # mean HF power across channels, per sample + # Cap the smoothing window at the signal length: np.convolve(mode="same") + # returns length max(len(env), w), so a window longer than the signal would + # otherwise yield a mask longer than n_times (breaks short recordings/epochs). + w = max(1, min(int(smooth_s * float(sfreq)), env.size)) + env = np.convolve(env, np.ones(w) / w, mode="same") + return env > np.quantile(env, float(quantile)) + + +def mwf_filter(X: np.ndarray, mask: np.ndarray, reg: float = 1e-6) -> np.ndarray: + """Multi-channel Wiener filter clean estimate. + + Parameters + ---------- + X : ndarray, shape (n_channels, n_times) + Multichannel signal. + mask : ndarray of bool, shape (n_times,) + ``True`` = artifact segment, ``False`` = clean segment. + reg : float + Diagonal-loading factor for covariance invertibility. + + Returns + ------- + cleaned : ndarray, shape (n_channels, n_times) + Wiener-filtered signal. + """ + X = np.asarray(X, dtype=float) + mask = np.asarray(mask, dtype=bool) # tolerate int 0/1 masks (avoid ~int bitwise NOT) + M = X.shape[0] + Xa, Xc = X[:, mask], X[:, ~mask] + if M < 2 or Xa.shape[1] < M + 1 or Xc.shape[1] < M + 1: + # spatial filter is degenerate for 1 channel, or too few samples in one + # segment to estimate covariances + return X.copy() + Ryy = np.cov(Xa) + Rnn = np.cov(Xc) + Ryy = Ryy + reg * (np.trace(Ryy) / M) * np.eye(M) # diagonal loading + return Rnn @ np.linalg.solve(Ryy, X) # R_clean @ R_artifact^{-1} @ X + + +def compute_mwf( + X: np.ndarray, + sfreq: float, + hf_hz: float = 20.0, + quantile: float = 0.6, + reg: float = 1e-6, + mask: np.ndarray | None = None, +) -> tuple[np.ndarray, dict[str, Any]]: + """Multi-channel Wiener filter cleaning of a data array. + + Parameters + ---------- + X : ndarray, shape (n_channels, n_times) + Multichannel signal. + sfreq : float + Sampling frequency in Hz (used only when ``mask`` is None). + hf_hz, quantile + Passed to :func:`hf_power_mask` when ``mask`` is None. + reg : float + Diagonal-loading factor for covariance invertibility. + mask : ndarray of bool | None + Optional artifact-segment mask. If None, it is estimated from broadband + high-frequency power. + + Returns + ------- + X_clean : ndarray, shape (n_channels, n_times) + Cleaned signal. + info : dict + Diagnostics: ``mask`` (the artifact mask used) and + ``artifact_fraction`` (fraction of samples flagged). + """ + X = np.asarray(X, dtype=float) + if X.ndim != 2: + raise ValueError( + f"Expected a 2-D (n_channels, n_samples) array, got shape {X.shape}." + ) + if mask is None: + mask = hf_power_mask(X, sfreq, hf_hz, quantile) + mask = np.asarray(mask, dtype=bool) # tolerate int 0/1 masks + if mask.all() or (~mask).all(): # no contrast -> nothing to estimate + cleaned = X.copy() + else: + cleaned = mwf_filter(X, mask, reg) + info = {"mask": mask, "artifact_fraction": float(np.mean(mask))} + return cleaned, info + + +class MWF(BaseEstimator, TransformerMixin): + """Multi-channel Wiener filter artifact remover (Somers et al. 2018). + + A generic, reference-free spatial cleaner (the RELAX-pipeline core). ``fit`` + estimates the artifact/clean covariances and the resulting Wiener operator on + the training data; ``transform`` applies that fixed operator to new data + (leakage-safe). Artifact segments are found from broadband high-frequency + power unless a mask is supplied. Accepts MNE ``Raw``/``Epochs`` objects or + NumPy ``(n_channels, n_samples)`` arrays; for MNE objects the sampling + frequency is read from ``info`` when ``sfreq`` is not given. + + .. note:: + + MWF is a general cleaner, not an artifact-specific method: its HF-power + segment split can attenuate genuine neural high-frequency activity when + the clean/artifact covariances are poorly separated. Validate preservation + of the band of interest on your data. + + Parameters + ---------- + sfreq : float | None + Sampling frequency in Hz. May be omitted for MNE input (read from + ``info['sfreq']``); required for NumPy-array input when no ``mask`` is + passed to ``transform``. + hf_hz : float + High-pass cutoff (Hz) for the high-frequency artifact detector. + quantile : float + HF-power quantile above which samples are flagged as artifact. + reg : float + Diagonal-loading factor for covariance invertibility. + verbose : bool | str | int | None + Control logging verbosity (MNE-style). + + Attributes + ---------- + spatial_filter_ : ndarray, shape (n_channels, n_channels) + The learned Wiener operator ``R_clean @ R_artifact^{-1}``. + artifact_mask_ : ndarray of bool, shape (n_train_times,) + Artifact-segment mask used to fit the operator. + artifact_fraction_ : float + Fraction of training samples flagged as artifact. + + Examples + -------- + >>> import numpy as np + >>> from mne_denoise.mwf import MWF + >>> rng = np.random.default_rng(0) + >>> X = rng.standard_normal((16, 4000)) + >>> cleaned = MWF(sfreq=250.0).fit_transform(X) + >>> cleaned.shape + (16, 4000) + """ + + def __init__( + self, + sfreq: float | None = None, + hf_hz: float = 20.0, + quantile: float = 0.6, + reg: float = 1e-6, + verbose: bool | str | int | None = None, + ) -> None: + self.sfreq = sfreq + self.hf_hz = hf_hz + self.quantile = quantile + self.reg = reg + self.verbose = verbose + + def _resolve_sfreq(self, sfreq_data: float | None) -> float | None: + return sfreq_data if sfreq_data is not None else self.sfreq + + def _to_2d(self, X: Any) -> tuple[np.ndarray, float | None, str, Any, Any]: + data, sfreq_data, mne_type, orig_inst, picks, _names = extract_data_from_mne( + X, auto_pick=True + ) + sfreq = self._resolve_sfreq(sfreq_data) + if mne_type == "epochs": + n_ep, n_ch, n_t = data.shape + data2d = np.transpose(data, (1, 0, 2)).reshape(n_ch, n_ep * n_t) + else: + data2d = np.asarray(data, dtype=float) + return data2d, sfreq, mne_type, orig_inst, picks + + def fit(self, X: Any, y=None, mask: np.ndarray | None = None) -> "MWF": + """Estimate the Wiener operator on the training data. + + Parameters + ---------- + X : Raw | Epochs | ndarray + Training data. Epochs are concatenated along time. + y : None + Ignored. + mask : ndarray of bool | None + Optional artifact-segment mask over the (concatenated) training + samples. If None, it is estimated from broadband HF power. + + Returns + ------- + self : MWF + """ + data2d, sfreq, _mne_type, _orig, _picks = self._to_2d(X) + M = data2d.shape[0] + if mask is None: + if sfreq is None: + raise ValueError( + "sfreq is required to estimate the artifact mask for array " + "input (pass MWF(sfreq=...)) or supply an explicit mask, or " + "fit on an MNE object carrying a sampling frequency." + ) + mask = hf_power_mask(data2d, sfreq, self.hf_hz, self.quantile) + mask = np.asarray(mask, dtype=bool) # tolerate int 0/1 masks + self.artifact_mask_ = mask + self.artifact_fraction_ = float(np.mean(mask)) + + Xa, Xc = data2d[:, mask], data2d[:, ~mask] + if ( + M < 2 + or mask.all() + or (~mask).all() + or Xa.shape[1] < M + 1 + or Xc.shape[1] < M + 1 + ): + # 1 channel is degenerate for a spatial filter, or insufficient + # contrast -> identity operator (no cleaning) + self.spatial_filter_ = np.eye(M) + else: + Ryy = np.cov(Xa) + Rnn = np.cov(Xc) + Ryy = Ryy + self.reg * (np.trace(Ryy) / M) * np.eye(M) + self.spatial_filter_ = Rnn @ np.linalg.inv(Ryy) + if self.verbose: + logger.info( + "MWF: fit on %.1f%% artifact samples.", 100.0 * self.artifact_fraction_ + ) + return self + + def transform(self, X: Any, y=None) -> Any: + """Apply the learned Wiener operator. + + Parameters + ---------- + X : Raw | Epochs | ndarray + Data to clean (same channel layout as the fitted data). + y : None + Ignored. + + Returns + ------- + X_clean : Raw | Epochs | ndarray + Cleaned data in the same format as the input. + """ + check_is_fitted(self, "spatial_filter_") + data, _sfreq, mne_type, orig_inst, picks, _names = extract_data_from_mne( + X, auto_pick=True + ) + if mne_type == "epochs": + cleaned = np.empty_like(data, dtype=float) + for e in range(data.shape[0]): + cleaned[e] = self.spatial_filter_ @ np.asarray(data[e], dtype=float) + else: + cleaned = self.spatial_filter_ @ np.asarray(data, dtype=float) + return reconstruct_mne_object( + cleaned, orig_inst, mne_type, picks=picks, verbose=False + ) + + def fit_transform(self, X: Any, y=None, mask: np.ndarray | None = None, **fit_params) -> Any: + """Fit on ``X`` and apply to ``X`` in one step. + + Parameters + ---------- + X : Raw | Epochs | ndarray + Input data. + y : None + Ignored. + mask : ndarray of bool | None + Optional artifact-segment mask (see :meth:`fit`). + **fit_params + Ignored. + + Returns + ------- + X_clean : Raw | Epochs | ndarray + Cleaned data. + """ + return self.fit(X, y, mask=mask).transform(X) + diff --git a/tests/test_mwf.py b/tests/test_mwf.py new file mode 100644 index 00000000..c2715b2f --- /dev/null +++ b/tests/test_mwf.py @@ -0,0 +1,218 @@ +"""Tests for the mne_denoise.mwf module (multi-channel Wiener filter).""" + +from __future__ import annotations + +import numpy as np +import pytest + +from mne_denoise.mwf import MWF, compute_mwf, hf_power_mask, mwf_filter + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def rng(): + """Shared random generator.""" + return np.random.default_rng(0) + + +@pytest.fixture() +def burst_data(rng): + """Synthetic EEG with intermittent broadband high-frequency bursts. + + Returns ``(X, sfreq, burst)`` where ``burst`` is the boolean sample mask of + the artifact segments. The clean background is low-frequency; the artifact is + a broadband HF burst injected into a subset of samples. + """ + sfreq = 250.0 + n_times = 4000 + n_ch = 12 + t = np.arange(n_times) / sfreq + + # Low-frequency neural background (well separated from the HF artifact). + neural = np.vstack( + [np.sin(2 * np.pi * f * t + rng.uniform(0, 2 * np.pi)) for f in (6.0, 8.0, 10.0)] + ) + M_neural = rng.standard_normal((n_ch, 3)) + X = M_neural @ neural + + # Broadband HF bursts on ~30% of samples. + burst = np.zeros(n_times, dtype=bool) + for start in (400, 1500, 2600, 3300): + burst[start : start + 300] = True + hf = rng.standard_normal((n_ch, n_times)) * 3.0 + from scipy.signal import butter, filtfilt + + b, a = butter(4, 40.0 / (0.5 * sfreq), btype="high") + hf = filtfilt(b, a, hf, axis=-1) + X = X + hf * burst + return X, sfreq, burst + + +def _band_power(X, sfreq, fmin, fmax): + spec = np.abs(np.fft.rfft(np.atleast_2d(X), axis=-1)) ** 2 + freqs = np.fft.rfftfreq(X.shape[-1], 1.0 / sfreq) + band = (freqs >= fmin) & (freqs < fmax) + return float(spec[:, band].sum()) + + +# --------------------------------------------------------------------------- +# Functional API +# --------------------------------------------------------------------------- + + +def test_hf_power_mask_flags_bursts(burst_data): + """The HF-power mask overlaps the injected burst samples.""" + X, sfreq, burst = burst_data + mask = hf_power_mask(X, sfreq, hf_hz=20.0, quantile=0.6) + assert mask.shape == (X.shape[-1],) + # Detected artifact samples are enriched inside the true bursts. + overlap = (mask & burst).sum() / max(mask.sum(), 1) + assert overlap > 0.5 + + +def test_compute_mwf_shapes_and_info(burst_data): + """compute_mwf returns cleaned data + a diagnostics dict.""" + X, sfreq, _burst = burst_data + cleaned, info = compute_mwf(X, sfreq) + assert cleaned.shape == X.shape + assert info["mask"].shape == (X.shape[-1],) + assert 0.0 <= info["artifact_fraction"] <= 1.0 + + +def test_compute_mwf_rejects_1d(): + """A 1-D input raises a clear error.""" + with pytest.raises(ValueError, match="2-D"): + compute_mwf(np.zeros(100), 250.0) + + +def test_mwf_filter_identity_when_covariances_match(rng): + """With a random mask on stationary data the filter is ~identity (no over-clean).""" + X = rng.standard_normal((8, 4000)) + mask = np.zeros(4000, dtype=bool) + mask[::2] = True # arbitrary split of stationary data + cleaned = mwf_filter(X, mask, reg=1e-6) + # Stationary data -> R_clean ~ R_artifact -> filter ~ identity. + assert np.corrcoef(cleaned.ravel(), X.ravel())[0, 1] > 0.9 + + +# --------------------------------------------------------------------------- +# MWF estimator +# --------------------------------------------------------------------------- + + +def test_mwf_fit_transform_numpy_shape(burst_data): + """fit_transform on a NumPy array returns an array of the same shape.""" + X, sfreq, _burst = burst_data + cleaned = MWF(sfreq=sfreq).fit_transform(X) + assert isinstance(cleaned, np.ndarray) + assert cleaned.shape == X.shape + + +def test_mwf_reduces_burst_hf_power(burst_data): + """MWF attenuates the broadband HF burst power.""" + X, sfreq, _burst = burst_data + cleaned = MWF(sfreq=sfreq, hf_hz=20.0, quantile=0.6).fit_transform(X) + hf_before = _band_power(X, sfreq, 40.0, 120.0) + hf_after = _band_power(cleaned, sfreq, 40.0, 120.0) + assert hf_after < hf_before + + +def test_mwf_fitted_attributes(burst_data): + """Fitted attributes are populated with correct shapes.""" + X, sfreq, _burst = burst_data + est = MWF(sfreq=sfreq).fit(X) + assert est.spatial_filter_.shape == (X.shape[0], X.shape[0]) + assert est.artifact_mask_.shape == (X.shape[-1],) + assert 0.0 <= est.artifact_fraction_ <= 1.0 + + +def test_mwf_leakage_split_applies_operator(rng): + """transform applies the operator learned in fit (train != eval).""" + train = rng.standard_normal((8, 4000)) + evalu = rng.standard_normal((8, 1000)) + mask = np.zeros(4000, dtype=bool) + mask[:1500] = True + est = MWF().fit(train, mask=mask) # explicit mask -> no sfreq needed + cleaned = est.transform(evalu) + np.testing.assert_allclose(cleaned, est.spatial_filter_ @ evalu) + + +def test_mwf_requires_sfreq_for_array_without_mask(burst_data): + """Array input without sfreq or mask raises a clear error.""" + X, _sfreq, _burst = burst_data + with pytest.raises(ValueError, match="sfreq is required"): + MWF().fit(X) + + +def test_mwf_transform_before_fit_raises(rng): + """transform before fit raises NotFittedError.""" + from sklearn.exceptions import NotFittedError + + with pytest.raises(NotFittedError): + MWF().transform(rng.standard_normal((8, 100))) + + +# --------------------------------------------------------------------------- +# MNE round-trip +# --------------------------------------------------------------------------- + + +def test_mwf_mne_raw_roundtrip_infers_sfreq(burst_data): + """fit_transform on an MNE Raw returns a Raw of identical shape; sfreq inferred.""" + mne = pytest.importorskip("mne") + X, sfreq, _burst = burst_data + info = mne.create_info([f"EEG{i:02d}" for i in range(X.shape[0])], sfreq, "eeg") + raw = mne.io.RawArray(X, info, verbose=False) + + cleaned = MWF(hf_hz=20.0, quantile=0.6).fit_transform(raw) # sfreq from info + assert isinstance(cleaned, mne.io.BaseRaw) + assert cleaned.get_data().shape == X.shape + assert not np.allclose(cleaned.get_data(), X) + + +# --------------------------------------------------------------------------- +# Edge cases +# --------------------------------------------------------------------------- + + +def test_hf_power_mask_short_signal_length(rng): + """The mask length equals n_times even when the smoothing window is longer.""" + # 20 samples at 250 Hz -> default smooth window (25 samples) exceeds length. + X = rng.standard_normal((8, 20)) + mask = hf_power_mask(X, 250.0) + assert mask.shape == (20,) + + +def test_mwf_short_signal_no_crash(rng): + """compute_mwf / MWF.fit_transform do not crash on very short signals.""" + X = rng.standard_normal((8, 20)) + cleaned, info = compute_mwf(X, 250.0) + assert cleaned.shape == X.shape + assert info["mask"].shape == (20,) + cleaned2 = MWF(sfreq=250.0).fit_transform(X) + assert cleaned2.shape == X.shape + + +def test_mwf_single_channel_is_identity(rng): + """A single channel is degenerate for a spatial filter -> returned unchanged.""" + X = rng.standard_normal((1, 2000)) + cleaned = MWF(sfreq=250.0).fit_transform(X) + assert cleaned.shape == X.shape + np.testing.assert_allclose(cleaned, X) + # functional API too + cleaned_f, _ = compute_mwf(X, 250.0) + assert cleaned_f.shape == X.shape + + +def test_mwf_int_mask_matches_bool_mask(rng): + """An integer 0/1 mask gives the same result as the equivalent bool mask.""" + X = rng.standard_normal((8, 2000)) + bool_mask = np.zeros(2000, dtype=bool) + bool_mask[:800] = True + int_mask = bool_mask.astype(int) + cleaned_bool, _ = compute_mwf(X, 250.0, mask=bool_mask) + cleaned_int, _ = compute_mwf(X, 250.0, mask=int_mask) + np.testing.assert_allclose(cleaned_bool, cleaned_int) From 1b129e1a6e9765d0126f42614c29c44e1d04816d Mon Sep 17 00:00:00 2001 From: Sina Esmaeili Date: Fri, 31 Jul 2026 22:40:54 -0400 Subject: [PATCH 2/4] FIX: align MWF with zero-delay GEVD contracts --- docs/api.rst | 7 +- docs/changes/devel/51.feature.rst | 3 + docs/mwf.md | 133 ++-- mne_denoise/dss/denoisers/temporal.py | 2 +- mne_denoise/mwf/__init__.py | 34 +- mne_denoise/mwf/core.py | 951 +++++++++++++++++++------- tests/test_mwf.py | 546 +++++++++------ 7 files changed, 1128 insertions(+), 548 deletions(-) create mode 100644 docs/changes/devel/51.feature.rst diff --git a/docs/api.rst b/docs/api.rst index 046b9812..426d6ccf 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -63,17 +63,20 @@ iCanClean mne_denoise.icanclean.ICanClean mne_denoise.icanclean.compute_icanclean -MWF ---- +Multi-channel Wiener filter +--------------------------- .. autosummary:: :toctree: generated/ :nosignatures: + mne_denoise.mwf.MultichannelWienerFilter mne_denoise.mwf.MWF mne_denoise.mwf.compute_mwf mne_denoise.mwf.hf_power_mask mne_denoise.mwf.mwf_filter +``MWF`` is a short alias of ``MultichannelWienerFilter``. + Denoisers --------- .. autosummary:: diff --git a/docs/changes/devel/51.feature.rst b/docs/changes/devel/51.feature.rst new file mode 100644 index 00000000..379cdeae --- /dev/null +++ b/docs/changes/devel/51.feature.rst @@ -0,0 +1,3 @@ +Added the zero-delay GEVD :class:`mne_denoise.mwf.MultichannelWienerFilter` +with explicit artifact-mask and clean-reference contracts, diagnostics, and MNE +metadata preservation. ``MWF`` remains available as a documented short alias. diff --git a/docs/mwf.md b/docs/mwf.md index 743b8f19..ba4f793d 100644 --- a/docs/mwf.md +++ b/docs/mwf.md @@ -1,83 +1,108 @@ -# Multi-channel Wiener Filter (MWF) +# Multi-channel Wiener filtering -## Overview +`mne_denoise.mwf.MultichannelWienerFilter` implements the zero-delay GEVD +multi-channel Wiener filter described by Somers, Francart, and Bertrand (2018). +`MWF` is a short compatibility alias of the same class. -The `mne_denoise.mwf` module implements the **multi-channel Wiener filter (MWF)** -(Somers, Francart & Bertrand 2018), a generic, reference-free spatial artifact -cleaner and the spatial-filter core of the **RELAX** pipeline (Bailey et al. -2023). +MWF is semi-supervised. It needs examples of both artifact-present and clean EEG +to estimate their covariance matrices. The core algorithm does not discover +artifacts automatically, and the origin of the mask is part of the scientific +operating point. -With the recording split into artifact-present and artifact-free segments, the -clean signal is recovered by +## Explicit-mask workflow +Mask values have precise meanings: + +- `1`: artifact-present training sample; +- `0`: clean training sample; +- `NaN`: ignored, or reassigned through `treat_nan`. + +```python +from mne_denoise.mwf import MultichannelWienerFilter + +mwf = MultichannelWienerFilter(rank="positive") +mwf.fit(train_raw, artifact_mask=train_mask) +cleaned = mwf.transform(eval_raw) ``` -X_clean = R_clean · R_artifact⁻¹ · X -``` -where `R_artifact` is the covariance over artifact segments (signal + artifact) -and `R_clean` the covariance over artifact-free segments (signal only). During -clean segments the two covariances coincide, so the filter is ~identity (no -over-cleaning); during artifact segments it projects out the artifact subspace. -No reference channel is required — artifact segments are marked from broadband -high-frequency power (or a caller-supplied mask). +`transform()` uses only the frozen spatial operator. It does not inspect the +evaluation recording or create a new mask. + +For epoched input, a mask can have shape `(n_epochs, n_times)` or be a flat +epoch-major vector. MNE channel names are aligned at transform time, and channels +not selected for fitting are preserved unchanged. -> **Note.** MWF is a *general* cleaner, not an artifact-specific method. Because -> its clean/artifact split is driven by broadband HF power, it can attenuate -> genuine neural high-frequency activity when the two covariances are poorly -> separated. Validate preservation of the band of interest on your data. It is -> provided here as the well-cited RELAX-core building block. +## Independent clean reference -## Quick Start +A separate clean recording can supply the clean covariance. If no mask is +provided, all samples in `X` train the artifact-present covariance: ```python -import numpy as np -from mne_denoise.mwf import MWF +mwf.fit(artifact_training_raw, clean_reference=clean_reference_raw) +``` -data = np.random.randn(32, 20000) # 32 channels, 20000 samples +The training and reference data must have the same channels, channel scaling, +physical units, and—when both are MNE objects—the same sampling frequency. -# Leakage-safe estimator API: learn the operator on train, apply to eval. -est = MWF(sfreq=250.0) -est.fit(data) -cleaned = est.transform(data) -print(f"fit on {est.artifact_fraction_:.0%} artifact samples") -``` +## Optional high-frequency mask authoring -Supplying an explicit artifact mask (no HF detector, no `sfreq` needed): +`hf_power_mask()` is a convenience heuristic, not part of the reference MWF +algorithm. It can be used explicitly: ```python -cleaned = MWF().fit_transform(data, mask=my_artifact_mask) +from mne_denoise.mwf import hf_power_mask + +mask = hf_power_mask(train_data, sfreq=250.0, hf_hz=20.0, quantile=0.7) +mwf.fit(train_data, artifact_mask=mask) ``` -With MNE-Python objects the sampling frequency is read from `info`: +Or requested as an explicit estimator strategy: ```python -raw_clean = MWF().fit_transform(raw) +mwf = MultichannelWienerFilter( + mask_strategy="hf_power", + sfreq=250.0, + hf_hz=20.0, + quantile=0.7, +) +mwf.fit(train_data) ``` -## One-shot functional API +This detector can label genuine high-frequency neural activity as artifact. Its +cutoff, quantile, and smoothing duration must therefore be validated for the +acquisition regime. -```python -from mne_denoise.mwf import compute_mwf, hf_power_mask, mwf_filter +## GEVD rank and diagnostics -cleaned, info = compute_mwf(data, sfreq=250.0) # (cleaned, {mask, artifact_fraction}) -mask = hf_power_mask(data, sfreq=250.0) # broadband HF artifact detector -cleaned = mwf_filter(data, mask) # raw Wiener filter given a mask -``` +The default `rank="positive"` matches the reference MATLAB toolbox's `poseig` +setting: only positive artifact eigenvalues are retained. `rank="full"` applies +the full-rank covariance-ratio filter, and an integer retains that many leading +GEVD directions. + +After fitting, the estimator exposes: + +- `generalized_eigenvalues_` and `artifact_eigenvalues_`; +- `selected_components_`; +- `artifact_mask_` and `artifact_fraction_`; +- `fit_diagnostics_`, including sample counts, covariance ranks, and the actual + diagonal loading. + +Relative diagonal loading makes the operator invariant to a shared global unit +rescaling. It cannot correct channel-specific unit mismatches. -## Parameters +## Evidence boundary -| Parameter | Description | -| ---------- | --------------------------------------------------------------------- | -| `sfreq` | Sampling frequency (Hz). Optional for MNE input / when a mask is given. | -| `hf_hz` | High-pass cutoff (Hz) for the HF artifact detector. | -| `quantile` | HF-power quantile above which samples are flagged as artifact. | -| `reg` | Diagonal-loading factor for covariance invertibility. | +This implementation is derived from the authors' public MATLAB equations and +locks internal invariants such as full-rank covariance-ratio equivalence. It does +not yet claim external numerical parity with the MATLAB toolbox or validated +performance for a particular acquisition regime. The reference toolbox also +supports temporal delay embedding; this implementation currently covers the +zero-delay method only. ## References 1. Somers, B., Francart, T., & Bertrand, A. (2018). A generic EEG artifact - removal algorithm based on the multi-channel Wiener filter. _Journal of Neural - Engineering_, 15(3), 036007. https://doi.org/10.1088/1741-2552/aaac92 -2. Bailey, N. W., et al. (2023). RELAX: An automated pre-processing pipeline for - cleaning EEG data — Part 1. _Clinical Neurophysiology_, 149, 178-201. - https://doi.org/10.1016/j.clinph.2023.01.007 + removal algorithm based on the multi-channel Wiener filter. *Journal of + Neural Engineering*, 15(3), 036007. https://doi.org/10.1088/1741-2552/aaac92 +2. Authors' MATLAB implementation: + https://github.com/exporl/mwf-artifact-removal diff --git a/mne_denoise/dss/denoisers/temporal.py b/mne_denoise/dss/denoisers/temporal.py index 7b5bf056..0623cbc7 100644 --- a/mne_denoise/dss/denoisers/temporal.py +++ b/mne_denoise/dss/denoisers/temporal.py @@ -158,7 +158,7 @@ class SmoothingBias(LinearDenoiser): Uses a boxcar moving average filter to smooth the data. When used to split the signal into a smooth branch and a residual (``data - smooth``), fitting DSS on the residual and adding the smooth branch back follows ZapLine's - period-matched decomposition (de Cheveigné, 2020 [3]_): with + period-matched decomposition (de Cheveigné, 2020): with ``window = round(sfreq / f_line)`` the smoother has zeros at ``f_line`` and its harmonics, so the residual concentrates the narrowband artifact. diff --git a/mne_denoise/mwf/__init__.py b/mne_denoise/mwf/__init__.py index d6b7034d..3e1e1b2e 100644 --- a/mne_denoise/mwf/__init__.py +++ b/mne_denoise/mwf/__init__.py @@ -1,30 +1,22 @@ -"""Multi-channel Wiener filter (MWF) artifact removal. +"""Semi-supervised multi-channel Wiener filtering. -This module contains: - -- ``hf_power_mask``: broadband high-frequency artifact-segment detector. -- ``mwf_filter`` / ``compute_mwf``: the array-based Wiener filter. -- ``MWF``: the scikit-learn estimator, compatible with MNE-Python objects or - NumPy arrays. - -The MWF (Somers, Francart & Bertrand 2018) is a generic, reference-free spatial -artifact cleaner and the spatial-filter core of the RELAX pipeline. It recovers -the clean signal via ``R_clean @ R_artifact^{-1} @ X`` from artifact-free and -artifact-present segment covariances. It is a general cleaner rather than an -artifact-specific method; validate preservation of the band of interest. - -References ----------- -.. [1] Somers, B., Francart, T., & Bertrand, A. (2018). A generic EEG artifact - removal algorithm based on the multi-channel Wiener filter. Journal of - Neural Engineering, 15(3), 036007. - https://doi.org/10.1088/1741-2552/aaac92 +``MultichannelWienerFilter`` is the canonical estimator name. ``MWF`` is a +documented short alias. The implementation uses the zero-delay GEVD formulation +from Somers, Francart, and Bertrand (2018); explicit temporal delay embedding is +not yet implemented. """ -from .core import MWF, compute_mwf, hf_power_mask, mwf_filter +from .core import ( + MWF, + MultichannelWienerFilter, + compute_mwf, + hf_power_mask, + mwf_filter, +) __all__ = [ "MWF", + "MultichannelWienerFilter", "compute_mwf", "hf_power_mask", "mwf_filter", diff --git a/mne_denoise/mwf/core.py b/mne_denoise/mwf/core.py index dc9caea4..d7ff0701 100644 --- a/mne_denoise/mwf/core.py +++ b/mne_denoise/mwf/core.py @@ -1,363 +1,790 @@ -"""Multi-channel Wiener filter (MWF) for EEG artifact removal. +"""GEVD multi-channel Wiener filtering for semi-supervised EEG denoising. -The multi-channel Wiener filter (Somers, Francart & Bertrand 2018) [1]_ is a -generic, reference-free artifact remover and the spatial-filter core of the -RELAX pipeline (Bailey et al. 2023) [2]_. With the recording split into -artifact-present and artifact-free segments, the clean signal is recovered by +This module implements the zero-delay form of the multi-channel Wiener filter +(MWF) described by Somers, Francart, and Bertrand (2018). The filter learns an +artifact covariance from explicitly marked artifact samples and a clean +covariance from marked clean samples or a separate clean reference. A +generalized eigendecomposition (GEVD) then constructs a low-rank artifact model. - X_clean = R_clean @ R_artifact^{-1} @ X - -where ``R_artifact`` is the covariance over artifact segments (signal + artifact) -and ``R_clean`` the covariance over artifact-free segments (signal only). During -clean segments the two covariances coincide so the filter is ~identity (no -over-cleaning); during artifact segments it projects out the artifact subspace. -No reference channel is needed: artifact segments are marked by a -high-frequency-power threshold (a caller-supplied mask is honoured). - -.. note:: - - MWF is a *general* spatial cleaner rather than an artifact-specific method. - Because its clean/artifact split is driven by broadband high-frequency power, - it can attenuate genuine neural high-frequency activity when the two - covariances are poorly separated; validate preservation of the band of - interest on your data. In our M/EEG benchmark it was less selective than the - reference-free CCA/SSA methods for muscle removal — it is provided here as the - well-cited RELAX-core building block, not as a targeted muscle remover. - -This module contains: - -- ``hf_power_mask``: broadband high-frequency artifact-segment detector. -- ``mwf_filter`` / ``compute_mwf``: the array-based Wiener filter. -- ``MWF``: the scikit-learn estimator, compatible with MNE-Python objects or - NumPy arrays. +The high-frequency detector provided here is only an optional mask-authoring +heuristic; it is not part of the core MWF algorithm and must be selected +explicitly. The authors' MATLAB toolbox also supports delay embedding, which is +outside the scope of this zero-delay implementation. References ---------- -.. [1] Somers, B., Francart, T., & Bertrand, A. (2018). A generic EEG artifact - removal algorithm based on the multi-channel Wiener filter. Journal of - Neural Engineering, 15(3), 036007. - https://doi.org/10.1088/1741-2552/aaac92 -.. [2] Bailey, N. W., et al. (2023). RELAX: An automated pre-processing pipeline - for cleaning EEG data - Part 1. Clinical Neurophysiology, 149, 178-201. - https://doi.org/10.1016/j.clinph.2023.01.007 +Somers, B., Francart, T., & Bertrand, A. (2018). A generic EEG artifact removal +algorithm based on the multi-channel Wiener filter. Journal of Neural +Engineering, 15(3), 036007. https://doi.org/10.1088/1741-2552/aaac92 """ from __future__ import annotations import logging -from typing import Any +from numbers import Integral, Real +from typing import Any, Literal import numpy as np +from scipy.linalg import LinAlgError, eigh +from scipy.signal import butter, sosfiltfilt from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted -from ..utils import extract_data_from_mne, reconstruct_mne_object +from ..utils import extract_data_from_mne logger = logging.getLogger(__name__) +RankSpec = Literal["positive", "full"] | int +MaskStrategy = Literal["hf_power"] | None + + +def _as_2d_finite( + X: np.ndarray, + *, + name: str, + min_channels: int = 1, + min_samples: int = 2, +) -> np.ndarray: + """Validate and return channel-by-sample floating-point data.""" + try: + data = np.asarray(X, dtype=float) + except (TypeError, ValueError) as err: + raise TypeError(f"{name} must be a numeric array.") from err + if data.ndim != 2: + raise ValueError( + f"{name} must have shape (n_channels, n_samples), got {data.shape}." + ) + if data.shape[0] < min_channels: + raise ValueError(f"{name} must contain at least {min_channels} channels.") + if data.shape[1] < min_samples: + raise ValueError(f"{name} must contain at least {min_samples} samples.") + if not np.all(np.isfinite(data)): + raise ValueError(f"{name} must contain only finite values.") + return data + + +def _validate_real( + value: float, + *, + name: str, + minimum: float | None = None, + strict: bool = False, +) -> float: + """Validate a finite real scalar.""" + if isinstance(value, bool) or not isinstance(value, Real) or not np.isfinite(value): + raise ValueError(f"{name} must be a finite real number.") + value = float(value) + if minimum is not None: + valid = value > minimum if strict else value >= minimum + if not valid: + relation = ">" if strict else ">=" + raise ValueError(f"{name} must be {relation} {minimum}.") + return value + + +def _normalize_mask( + artifact_mask: np.ndarray, + *, + n_samples: int, + treat_nan: Literal["ignore", "artifact", "clean"], +) -> np.ndarray: + """Return a flat float mask containing only zero, one, or NaN.""" + if treat_nan not in {"ignore", "artifact", "clean"}: + raise ValueError("treat_nan must be 'ignore', 'artifact', or 'clean'.") + mask = np.asarray(artifact_mask) + if mask.ndim not in (1, 2) or mask.size != n_samples: + raise ValueError( + "artifact_mask must be one-dimensional, or epoch-by-time, with " + f"exactly {n_samples} entries; got shape {mask.shape}." + ) + if mask.dtype == np.bool_: + normalized = mask.astype(float, copy=False).reshape(-1) + elif np.issubdtype(mask.dtype, np.number): + normalized = mask.astype(float, copy=False).reshape(-1) + finite = normalized[np.isfinite(normalized)] + if not np.all(np.isin(finite, (0.0, 1.0))): + raise ValueError("artifact_mask values must be 0, 1, or NaN.") + if np.any(np.isinf(normalized)): + raise ValueError("artifact_mask values must be 0, 1, or NaN.") + else: + raise TypeError("artifact_mask must contain boolean or numeric values.") + + if treat_nan == "artifact": + normalized = np.nan_to_num(normalized, nan=1.0) + elif treat_nan == "clean": + normalized = np.nan_to_num(normalized, nan=0.0) + return normalized + def hf_power_mask( X: np.ndarray, sfreq: float, + *, hf_hz: float = 20.0, quantile: float = 0.6, smooth_s: float = 0.1, ) -> np.ndarray: - """Mark high-frequency-power time points as artifact (broadband muscle/transient). + """Create an artifact mask from smoothed high-frequency power. + + This detector is a convenience heuristic, not part of the reference MWF + algorithm. Its operating point must be validated for the acquisition regime. Parameters ---------- - X : ndarray, shape (n_channels, n_times) - Multichannel signal. + X : ndarray, shape (n_channels, n_samples) + Multichannel signal in any consistent physical unit. sfreq : float Sampling frequency in Hz. hf_hz : float - High-pass cutoff (Hz) used to isolate broadband high-frequency power. + High-pass cutoff in Hz. It must lie strictly below Nyquist. quantile : float - Samples whose smoothed HF power exceeds this quantile are flagged. + Fractional power quantile strictly between zero and one. Samples above + this threshold are marked as artifact. smooth_s : float - Moving-average smoothing window (seconds) applied to the HF envelope. + Positive moving-average duration in seconds. Returns ------- - mask : ndarray of bool, shape (n_times,) - ``True`` where the sample is flagged as artifact. - """ - from scipy.signal import butter, filtfilt - - ny = 0.5 * float(sfreq) - b, a = butter(4, min(float(hf_hz) / ny, 0.99), btype="high") - hf = filtfilt(b, a, np.asarray(X, dtype=float), axis=-1) - env = (hf**2).mean(axis=0) # mean HF power across channels, per sample - # Cap the smoothing window at the signal length: np.convolve(mode="same") - # returns length max(len(env), w), so a window longer than the signal would - # otherwise yield a mask longer than n_times (breaks short recordings/epochs). - w = max(1, min(int(smooth_s * float(sfreq)), env.size)) - env = np.convolve(env, np.ones(w) / w, mode="same") - return env > np.quantile(env, float(quantile)) - - -def mwf_filter(X: np.ndarray, mask: np.ndarray, reg: float = 1e-6) -> np.ndarray: - """Multi-channel Wiener filter clean estimate. - - Parameters - ---------- - X : ndarray, shape (n_channels, n_times) - Multichannel signal. - mask : ndarray of bool, shape (n_times,) - ``True`` = artifact segment, ``False`` = clean segment. - reg : float - Diagonal-loading factor for covariance invertibility. + artifact_mask : ndarray of bool, shape (n_samples,) + True for samples selected as artifact-present. - Returns - ------- - cleaned : ndarray, shape (n_channels, n_times) - Wiener-filtered signal. + Notes + ----- + The output depends on ``sfreq``, ``hf_hz``, ``quantile``, and ``smooth_s``; + these values are part of the MWF operating point, not universal defaults. """ - X = np.asarray(X, dtype=float) - mask = np.asarray(mask, dtype=bool) # tolerate int 0/1 masks (avoid ~int bitwise NOT) - M = X.shape[0] - Xa, Xc = X[:, mask], X[:, ~mask] - if M < 2 or Xa.shape[1] < M + 1 or Xc.shape[1] < M + 1: - # spatial filter is degenerate for 1 channel, or too few samples in one - # segment to estimate covariances - return X.copy() - Ryy = np.cov(Xa) - Rnn = np.cov(Xc) - Ryy = Ryy + reg * (np.trace(Ryy) / M) * np.eye(M) # diagonal loading - return Rnn @ np.linalg.solve(Ryy, X) # R_clean @ R_artifact^{-1} @ X + data = _as_2d_finite(X, name="X") + sfreq = _validate_real(sfreq, name="sfreq", minimum=0.0, strict=True) + hf_hz = _validate_real(hf_hz, name="hf_hz", minimum=0.0, strict=True) + quantile = _validate_real(quantile, name="quantile", minimum=0.0, strict=True) + smooth_s = _validate_real(smooth_s, name="smooth_s", minimum=0.0, strict=True) + if hf_hz >= sfreq / 2.0: + raise ValueError("hf_hz must be strictly below the Nyquist frequency.") + if quantile >= 1.0: + raise ValueError("quantile must be strictly between 0 and 1.") + + sos = butter(4, hf_hz, btype="highpass", fs=sfreq, output="sos") + try: + high_frequency = sosfiltfilt(sos, data, axis=-1) + except ValueError as err: + raise ValueError( + "X is too short for zero-phase high-frequency mask estimation; " + "provide an explicit artifact_mask." + ) from err + envelope = np.mean(high_frequency**2, axis=0) + window = max(1, min(int(round(smooth_s * sfreq)), envelope.size)) + envelope = np.convolve(envelope, np.ones(window) / window, mode="same") + return envelope > np.quantile(envelope, quantile) + + +def _covariance(data: np.ndarray) -> np.ndarray: + """Return a symmetric sample covariance matrix.""" + covariance = np.atleast_2d(np.cov(data, rowvar=True, bias=False)) + return (covariance + covariance.T) / 2.0 + + +def _select_rank(rank: RankSpec, artifact_eigenvalues: np.ndarray) -> np.ndarray: + """Select generalized eigen-directions for the artifact model.""" + n_channels = artifact_eigenvalues.size + if isinstance(rank, str): + if rank == "full": + return np.arange(n_channels) + if rank != "positive": + raise ValueError("rank must be 'positive', 'full', or a positive integer.") + tolerance = ( + np.finfo(float).eps + * n_channels + * max(1.0, float(np.max(np.abs(artifact_eigenvalues)))) + ) + return np.flatnonzero(artifact_eigenvalues > tolerance) + if isinstance(rank, bool) or not isinstance(rank, Integral): + raise TypeError("rank must be 'positive', 'full', or a positive integer.") + if not 1 <= int(rank) <= n_channels: + raise ValueError(f"Integer rank must be between 1 and {n_channels}.") + return np.arange(int(rank)) + + +def _compute_operator( + artifact_data: np.ndarray, + clean_data: np.ndarray, + *, + rank: RankSpec, + artifact_weight: float, + reg: float, +) -> tuple[np.ndarray, dict[str, Any]]: + """Construct the zero-delay GEVD MWF cleaning operator.""" + artifact_weight = _validate_real( + artifact_weight, name="artifact_weight", minimum=0.0, strict=True + ) + reg = _validate_real(reg, name="reg", minimum=0.0) + if artifact_data.shape[0] != clean_data.shape[0]: + raise ValueError("Artifact and clean data must have the same channel count.") + if artifact_data.shape[1] < 2 or clean_data.shape[1] < 2: + raise ValueError( + "MWF needs at least two artifact and two clean training samples." + ) + + artifact_covariance = _covariance(artifact_data) + clean_covariance = _covariance(clean_data) + n_channels = artifact_covariance.shape[0] + artifact_rank = int(np.linalg.matrix_rank(artifact_covariance)) + clean_rank = int(np.linalg.matrix_rank(clean_covariance)) + covariance_scale = max( + float(np.trace(artifact_covariance) / n_channels), + float(np.trace(clean_covariance) / n_channels), + ) + if not np.isfinite(covariance_scale) or covariance_scale <= np.finfo(float).tiny: + raise ValueError("MWF training covariances have zero numerical energy.") + loading = reg * covariance_scale + identity = np.eye(n_channels) + artifact_covariance_loaded = artifact_covariance + loading * identity + clean_covariance_loaded = clean_covariance + loading * identity + + try: + generalized_eigenvalues, eigenvectors = eigh( + artifact_covariance_loaded, + clean_covariance_loaded, + check_finite=False, + ) + except LinAlgError as err: + raise ValueError( + "The clean covariance is not positive definite. Use a positive reg " + "or provide more independent clean samples." + ) from err + order = np.argsort(generalized_eigenvalues)[::-1] + generalized_eigenvalues = generalized_eigenvalues[order] + eigenvectors = eigenvectors[:, order] + artifact_eigenvalues = generalized_eigenvalues - 1.0 + selected = _select_rank(rank, artifact_eigenvalues) + + denominators = generalized_eigenvalues + artifact_weight - 1.0 + tolerance = np.finfo(float).eps * max(1.0, float(np.max(np.abs(denominators)))) + if selected.size and np.any(np.abs(denominators[selected]) <= tolerance): + raise ValueError( + "artifact_weight creates a singular MWF denominator for a selected " + "component." + ) + component_weights = np.zeros(n_channels) + component_weights[selected] = ( + artifact_eigenvalues[selected] / denominators[selected] + ) + weighted_vectors = eigenvectors * component_weights[None, :] + try: + reference_artifact_filter = np.linalg.solve( + eigenvectors.T, weighted_vectors.T + ).T + except np.linalg.LinAlgError as err: + raise ValueError("The generalized eigenvector matrix is singular.") from err + + # The reference MATLAB implementation applies W.T as the artifact operator. + artifact_operator = reference_artifact_filter.T + spatial_filter = identity - artifact_operator + if not np.all(np.isfinite(spatial_filter)): + raise ValueError("MWF produced a non-finite spatial filter.") + + diagnostics = { + "rank_requested": rank, + "rank_used": int(selected.size), + "selected_components": selected.copy(), + "generalized_eigenvalues": generalized_eigenvalues.copy(), + "artifact_eigenvalues": artifact_eigenvalues.copy(), + "artifact_covariance_rank": artifact_rank, + "clean_covariance_rank": clean_rank, + "regularization_loading": float(loading), + "used_identity": bool(selected.size == 0), + } + return spatial_filter, diagnostics + + +def _resolve_mask( + data: np.ndarray, + artifact_mask: np.ndarray | None, + *, + clean_reference: np.ndarray | None, + mask_strategy: MaskStrategy, + sfreq: float | None, + hf_hz: float, + quantile: float, + smooth_s: float, + treat_nan: Literal["ignore", "artifact", "clean"], +) -> np.ndarray: + """Resolve an explicit or explicitly requested heuristic mask.""" + if artifact_mask is not None and mask_strategy is not None: + raise ValueError("Pass artifact_mask or mask_strategy, not both.") + if artifact_mask is not None: + return _normalize_mask( + artifact_mask, n_samples=data.shape[1], treat_nan=treat_nan + ) + if mask_strategy is None and clean_reference is not None: + return np.ones(data.shape[1], dtype=float) + if mask_strategy is None: + raise ValueError( + "An explicit artifact_mask is required. To opt into the unvalidated " + "high-frequency heuristic, set mask_strategy='hf_power'." + ) + if mask_strategy != "hf_power": + raise ValueError("mask_strategy must be None or 'hf_power'.") + if sfreq is None: + raise ValueError("sfreq is required when mask_strategy='hf_power'.") + return hf_power_mask( + data, + sfreq, + hf_hz=hf_hz, + quantile=quantile, + smooth_s=smooth_s, + ).astype(float) + + +def _fit_from_training_data( + data: np.ndarray, + artifact_mask: np.ndarray | None, + *, + clean_reference: np.ndarray | None, + mask_strategy: MaskStrategy, + sfreq: float | None, + hf_hz: float, + quantile: float, + smooth_s: float, + rank: RankSpec, + artifact_weight: float, + reg: float, + treat_nan: Literal["ignore", "artifact", "clean"], +) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]: + """Resolve training segments and fit a spatial operator.""" + data = _as_2d_finite(data, name="X", min_channels=2) + if clean_reference is not None: + clean_reference = _as_2d_finite( + clean_reference, name="clean_reference", min_channels=2 + ) + if clean_reference.shape[0] != data.shape[0]: + raise ValueError( + "clean_reference must have the same number of channels as X." + ) + normalized_mask = _resolve_mask( + data, + artifact_mask, + clean_reference=clean_reference, + mask_strategy=mask_strategy, + sfreq=sfreq, + hf_hz=hf_hz, + quantile=quantile, + smooth_s=smooth_s, + treat_nan=treat_nan, + ) + artifact_samples = normalized_mask == 1.0 + clean_samples = normalized_mask == 0.0 + ignored_samples = np.isnan(normalized_mask) + if artifact_samples.sum() < 2: + raise ValueError("artifact_mask must select at least two artifact samples.") + if clean_reference is None: + if clean_samples.sum() < 2: + raise ValueError("artifact_mask must select at least two clean samples.") + clean_data = data[:, clean_samples] + else: + clean_data = clean_reference + spatial_filter, diagnostics = _compute_operator( + data[:, artifact_samples], + clean_data, + rank=rank, + artifact_weight=artifact_weight, + reg=reg, + ) + valid_samples = (~ignored_samples).sum() + diagnostics.update( + { + "artifact_samples": int(artifact_samples.sum()), + "clean_samples": int(clean_data.shape[1]), + "ignored_samples": int(ignored_samples.sum()), + "artifact_fraction": float(artifact_samples.sum() / max(1, valid_samples)), + "used_clean_reference": clean_reference is not None, + "mask_strategy": mask_strategy, + } + ) + return spatial_filter, normalized_mask, diagnostics + + +def _apply_spatial_filter(data: np.ndarray, spatial_filter: np.ndarray) -> np.ndarray: + """Apply a spatial filter after removing and then restoring channel means.""" + if data.ndim == 2: + channel_means = data.mean(axis=1, keepdims=True) + return spatial_filter @ (data - channel_means) + channel_means + if data.ndim == 3: + n_epochs, n_channels, n_times = data.shape + continuous = np.transpose(data, (1, 0, 2)).reshape(n_channels, -1) + channel_means = continuous.mean(axis=1, keepdims=True) + cleaned = spatial_filter @ (continuous - channel_means) + channel_means + return cleaned.reshape(n_channels, n_epochs, n_times).transpose(1, 0, 2) + raise ValueError(f"Data must be 2D or 3D, got shape {data.shape}.") def compute_mwf( X: np.ndarray, - sfreq: float, + artifact_mask: np.ndarray | None = None, + *, + clean_reference: np.ndarray | None = None, + mask_strategy: MaskStrategy = None, + sfreq: float | None = None, hf_hz: float = 20.0, quantile: float = 0.6, + smooth_s: float = 0.1, + rank: RankSpec = "positive", + artifact_weight: float = 1.0, reg: float = 1e-6, - mask: np.ndarray | None = None, + treat_nan: Literal["ignore", "artifact", "clean"] = "ignore", ) -> tuple[np.ndarray, dict[str, Any]]: - """Multi-channel Wiener filter cleaning of a data array. + """Fit and apply a zero-delay GEVD multi-channel Wiener filter. Parameters ---------- - X : ndarray, shape (n_channels, n_times) - Multichannel signal. - sfreq : float - Sampling frequency in Hz (used only when ``mask`` is None). - hf_hz, quantile - Passed to :func:`hf_power_mask` when ``mask`` is None. + X : ndarray, shape (n_channels, n_samples) + Training data and the data to clean. + artifact_mask : ndarray | None + Values 1 (artifact), 0 (clean), or NaN (handled by ``treat_nan``). + Required unless ``clean_reference`` contains the clean covariance data or + ``mask_strategy='hf_power'`` is selected explicitly. + clean_reference : ndarray | None + Optional clean data with the same channel order and physical units as + ``X``. When supplied, zero-valued mask samples in ``X`` are excluded and + this reference supplies the clean covariance. + mask_strategy : None | 'hf_power' + Optional explicit opt-in to :func:`hf_power_mask`. + sfreq, hf_hz, quantile, smooth_s + High-frequency mask parameters, used only by ``mask_strategy='hf_power'``. + rank : {'positive', 'full'} | int + GEVD artifact rank. ``'positive'`` matches the reference toolbox's + ``poseig`` default; an integer retains that many leading directions. + artifact_weight : float + Positive reference-toolbox noise weighting (``mu``); 1 is unweighted. reg : float - Diagonal-loading factor for covariance invertibility. - mask : ndarray of bool | None - Optional artifact-segment mask. If None, it is estimated from broadband - high-frequency power. + Non-negative relative diagonal loading applied to both covariances. + treat_nan : {'ignore', 'artifact', 'clean'} + Interpretation of NaNs in ``artifact_mask``. Returns ------- - X_clean : ndarray, shape (n_channels, n_times) - Cleaned signal. - info : dict - Diagnostics: ``mask`` (the artifact mask used) and - ``artifact_fraction`` (fraction of samples flagged). + cleaned : ndarray, shape (n_channels, n_samples) + Filtered data in the same physical units as ``X``. + diagnostics : dict + Mask, sample counts, covariance ranks, GEVD values, and rank selection. """ - X = np.asarray(X, dtype=float) - if X.ndim != 2: - raise ValueError( - f"Expected a 2-D (n_channels, n_samples) array, got shape {X.shape}." - ) - if mask is None: - mask = hf_power_mask(X, sfreq, hf_hz, quantile) - mask = np.asarray(mask, dtype=bool) # tolerate int 0/1 masks - if mask.all() or (~mask).all(): # no contrast -> nothing to estimate - cleaned = X.copy() - else: - cleaned = mwf_filter(X, mask, reg) - info = {"mask": mask, "artifact_fraction": float(np.mean(mask))} - return cleaned, info - - -class MWF(BaseEstimator, TransformerMixin): - """Multi-channel Wiener filter artifact remover (Somers et al. 2018). - - A generic, reference-free spatial cleaner (the RELAX-pipeline core). ``fit`` - estimates the artifact/clean covariances and the resulting Wiener operator on - the training data; ``transform`` applies that fixed operator to new data - (leakage-safe). Artifact segments are found from broadband high-frequency - power unless a mask is supplied. Accepts MNE ``Raw``/``Epochs`` objects or - NumPy ``(n_channels, n_samples)`` arrays; for MNE objects the sampling - frequency is read from ``info`` when ``sfreq`` is not given. - - .. note:: - - MWF is a general cleaner, not an artifact-specific method: its HF-power - segment split can attenuate genuine neural high-frequency activity when - the clean/artifact covariances are poorly separated. Validate preservation - of the band of interest on your data. + data = _as_2d_finite(X, name="X", min_channels=2) + spatial_filter, normalized_mask, diagnostics = _fit_from_training_data( + data, + artifact_mask, + clean_reference=clean_reference, + mask_strategy=mask_strategy, + sfreq=sfreq, + hf_hz=hf_hz, + quantile=quantile, + smooth_s=smooth_s, + rank=rank, + artifact_weight=artifact_weight, + reg=reg, + treat_nan=treat_nan, + ) + cleaned = _apply_spatial_filter(data, spatial_filter) + diagnostics = { + **diagnostics, + "artifact_mask": normalized_mask.copy(), + "spatial_filter": spatial_filter.copy(), + } + return cleaned, diagnostics + + +def mwf_filter( + X: np.ndarray, + artifact_mask: np.ndarray, + *, + clean_reference: np.ndarray | None = None, + rank: RankSpec = "positive", + artifact_weight: float = 1.0, + reg: float = 1e-6, + treat_nan: Literal["ignore", "artifact", "clean"] = "ignore", +) -> np.ndarray: + """Fit and apply MWF from an explicit artifact mask.""" + cleaned, _ = compute_mwf( + X, + artifact_mask, + clean_reference=clean_reference, + rank=rank, + artifact_weight=artifact_weight, + reg=reg, + treat_nan=treat_nan, + ) + return cleaned + + +def _reconstruct_like( + cleaned: np.ndarray, + orig_inst: Any, + mne_type: str, + picks: np.ndarray | None, +) -> Any: + """Insert cleaned channels into an exact copy of an MNE object.""" + if orig_inst is None or mne_type == "array": + return cleaned + output = orig_inst.copy() + selected = slice(None) if picks is None else picks + if mne_type == "raw": + output.load_data() + output._data[selected, :] = cleaned + elif mne_type == "epochs": + output.load_data() + output._data[:, selected, :] = cleaned + elif mne_type == "evoked": + output.data[selected, :] = cleaned + else: # pragma: no cover - guarded by extract_data_from_mne + raise TypeError(f"Unsupported MNE data type {mne_type!r}.") + return output + + +class MultichannelWienerFilter(BaseEstimator, TransformerMixin): + """Zero-delay GEVD multi-channel Wiener filter. + + The estimator is semi-supervised: ``fit`` requires an explicit artifact mask, + a clean reference, or an explicit opt-in to the high-frequency mask heuristic. + ``transform`` applies the frozen operator without fitting on evaluation data. Parameters ---------- - sfreq : float | None - Sampling frequency in Hz. May be omitted for MNE input (read from - ``info['sfreq']``); required for NumPy-array input when no ``mask`` is - passed to ``transform``. - hf_hz : float - High-pass cutoff (Hz) for the high-frequency artifact detector. - quantile : float - HF-power quantile above which samples are flagged as artifact. + rank : {'positive', 'full'} | int + GEVD artifact rank. ``'positive'`` retains positive artifact + eigenvalues, matching the reference toolbox default. + artifact_weight : float + Positive artifact/noise weighting parameter (reference ``mu``). reg : float - Diagonal-loading factor for covariance invertibility. + Non-negative relative covariance diagonal loading. + treat_nan : {'ignore', 'artifact', 'clean'} + Interpretation of NaNs in an explicit artifact mask. + mask_strategy : None | 'hf_power' + Mask-authoring strategy. The default requires explicit evidence; set + ``'hf_power'`` to opt into the heuristic detector. + sfreq : float | None + Sampling frequency in Hz for ``mask_strategy='hf_power'``. MNE metadata + is used when available and must agree with this value when both exist. + hf_hz, quantile, smooth_s : float + Operating point of the high-frequency mask heuristic. verbose : bool | str | int | None - Control logging verbosity (MNE-style). + Enable a concise fit summary when truthy. Attributes ---------- spatial_filter_ : ndarray, shape (n_channels, n_channels) - The learned Wiener operator ``R_clean @ R_artifact^{-1}``. - artifact_mask_ : ndarray of bool, shape (n_train_times,) - Artifact-segment mask used to fit the operator. - artifact_fraction_ : float - Fraction of training samples flagged as artifact. - - Examples - -------- - >>> import numpy as np - >>> from mne_denoise.mwf import MWF - >>> rng = np.random.default_rng(0) - >>> X = rng.standard_normal((16, 4000)) - >>> cleaned = MWF(sfreq=250.0).fit_transform(X) - >>> cleaned.shape - (16, 4000) + Frozen clean-signal spatial operator. + artifact_mask_ : ndarray, shape (n_training_samples,) + Normalized 0/1/NaN training mask. + generalized_eigenvalues_ : ndarray, shape (n_channels,) + Sorted GEVD eigenvalues. + artifact_eigenvalues_ : ndarray, shape (n_channels,) + GEVD values relative to the clean baseline (``lambda - 1``). + selected_components_ : ndarray + Artifact directions retained by ``rank``. + fit_diagnostics_ : dict + Sample counts, covariance ranks, and operating-point diagnostics. + + Notes + ----- + Inputs and a separate ``clean_reference`` must use the same channel scaling + and physical units. The relative regularization makes the operator invariant + to a common global rescaling, but not to channel-specific unit mismatches. + Delay embedding from the reference MATLAB toolbox is not implemented here. """ def __init__( self, + *, + rank: RankSpec = "positive", + artifact_weight: float = 1.0, + reg: float = 1e-6, + treat_nan: Literal["ignore", "artifact", "clean"] = "ignore", + mask_strategy: MaskStrategy = None, sfreq: float | None = None, hf_hz: float = 20.0, quantile: float = 0.6, - reg: float = 1e-6, + smooth_s: float = 0.1, verbose: bool | str | int | None = None, ) -> None: + self.rank = rank + self.artifact_weight = artifact_weight + self.reg = reg + self.treat_nan = treat_nan + self.mask_strategy = mask_strategy self.sfreq = sfreq self.hf_hz = hf_hz self.quantile = quantile - self.reg = reg + self.smooth_s = smooth_s self.verbose = verbose def _resolve_sfreq(self, sfreq_data: float | None) -> float | None: - return sfreq_data if sfreq_data is not None else self.sfreq - - def _to_2d(self, X: Any) -> tuple[np.ndarray, float | None, str, Any, Any]: - data, sfreq_data, mne_type, orig_inst, picks, _names = extract_data_from_mne( - X, auto_pick=True + """Resolve and cross-check configured and metadata sampling rates.""" + configured = self.sfreq + if configured is not None: + configured = _validate_real( + configured, name="sfreq", minimum=0.0, strict=True + ) + if sfreq_data is not None: + sfreq_data = _validate_real( + sfreq_data, name="data sfreq", minimum=0.0, strict=True + ) + if ( + configured is not None + and sfreq_data is not None + and not np.isclose(configured, sfreq_data, rtol=1e-9, atol=0.0) + ): + raise ValueError( + f"Configured sfreq ({configured}) does not match MNE metadata " + f"({sfreq_data})." + ) + return sfreq_data if sfreq_data is not None else configured + + @staticmethod + def _extract_fit_data( + X: Any, *, ch_names: list[str] | None = None + ) -> tuple[np.ndarray, float | None, list[str] | None]: + """Extract and concatenate training data.""" + data, sfreq, _, _, _, extracted_names = extract_data_from_mne( + X, + ch_names=ch_names, + auto_pick=True, + concatenate_epochs=True, ) - sfreq = self._resolve_sfreq(sfreq_data) - if mne_type == "epochs": - n_ep, n_ch, n_t = data.shape - data2d = np.transpose(data, (1, 0, 2)).reshape(n_ch, n_ep * n_t) - else: - data2d = np.asarray(data, dtype=float) - return data2d, sfreq, mne_type, orig_inst, picks + data = _as_2d_finite(data, name="X", min_channels=2) + return data, sfreq, extracted_names - def fit(self, X: Any, y=None, mask: np.ndarray | None = None) -> "MWF": - """Estimate the Wiener operator on the training data. + def fit( + self, + X: Any, + y=None, + *, + artifact_mask: np.ndarray | None = None, + clean_reference: Any | None = None, + ) -> MultichannelWienerFilter: + """Estimate the frozen MWF operator. Parameters ---------- - X : Raw | Epochs | ndarray - Training data. Epochs are concatenated along time. + X : Raw | Epochs | Evoked | ndarray + Training data. Epochs are concatenated in epoch-major time order. y : None Ignored. - mask : ndarray of bool | None - Optional artifact-segment mask over the (concatenated) training - samples. If None, it is estimated from broadband HF power. + artifact_mask : ndarray | None + Explicit 0/1/NaN mask. Epoch-shaped masks are flattened in the same + order as the training epochs. + clean_reference : Raw | Epochs | Evoked | ndarray | None + Optional independent clean data with matching channels and units. If + no mask is supplied, all samples in ``X`` train the artifact + covariance and all reference samples train the clean covariance. Returns ------- - self : MWF + self : MultichannelWienerFilter + Fitted estimator. """ - data2d, sfreq, _mne_type, _orig, _picks = self._to_2d(X) - M = data2d.shape[0] - if mask is None: - if sfreq is None: + del y + data, sfreq_data, ch_names = self._extract_fit_data(X) + sfreq = self._resolve_sfreq(sfreq_data) + reference_data = None + if clean_reference is not None: + reference_data, reference_sfreq, _ = self._extract_fit_data( + clean_reference, ch_names=ch_names + ) + if reference_data.shape[0] != data.shape[0]: + raise ValueError( + "clean_reference must have the same number of channels as X." + ) + if ( + sfreq_data is not None + and reference_sfreq is not None + and not np.isclose(sfreq_data, reference_sfreq, rtol=1e-9, atol=0.0) + ): raise ValueError( - "sfreq is required to estimate the artifact mask for array " - "input (pass MWF(sfreq=...)) or supply an explicit mask, or " - "fit on an MNE object carrying a sampling frequency." + "clean_reference sampling frequency must match the training data." ) - mask = hf_power_mask(data2d, sfreq, self.hf_hz, self.quantile) - mask = np.asarray(mask, dtype=bool) # tolerate int 0/1 masks - self.artifact_mask_ = mask - self.artifact_fraction_ = float(np.mean(mask)) - Xa, Xc = data2d[:, mask], data2d[:, ~mask] - if ( - M < 2 - or mask.all() - or (~mask).all() - or Xa.shape[1] < M + 1 - or Xc.shape[1] < M + 1 - ): - # 1 channel is degenerate for a spatial filter, or insufficient - # contrast -> identity operator (no cleaning) - self.spatial_filter_ = np.eye(M) - else: - Ryy = np.cov(Xa) - Rnn = np.cov(Xc) - Ryy = Ryy + self.reg * (np.trace(Ryy) / M) * np.eye(M) - self.spatial_filter_ = Rnn @ np.linalg.inv(Ryy) + spatial_filter, normalized_mask, diagnostics = _fit_from_training_data( + data, + artifact_mask, + clean_reference=reference_data, + mask_strategy=self.mask_strategy, + sfreq=sfreq, + hf_hz=self.hf_hz, + quantile=self.quantile, + smooth_s=self.smooth_s, + rank=self.rank, + artifact_weight=self.artifact_weight, + reg=self.reg, + treat_nan=self.treat_nan, + ) + self.spatial_filter_ = spatial_filter + self.artifact_mask_ = normalized_mask + self.artifact_fraction_ = diagnostics["artifact_fraction"] + self.generalized_eigenvalues_ = diagnostics["generalized_eigenvalues"] + self.artifact_eigenvalues_ = diagnostics["artifact_eigenvalues"] + self.selected_components_ = diagnostics["selected_components"] + self.fit_diagnostics_ = diagnostics + self._n_channels_ = data.shape[0] + self._fit_ch_names_ = ch_names if self.verbose: logger.info( - "MWF: fit on %.1f%% artifact samples.", 100.0 * self.artifact_fraction_ + "MWF fit: %d artifact samples, %d clean samples, rank %d.", + diagnostics["artifact_samples"], + diagnostics["clean_samples"], + diagnostics["rank_used"], ) return self - def transform(self, X: Any, y=None) -> Any: - """Apply the learned Wiener operator. - - Parameters - ---------- - X : Raw | Epochs | ndarray - Data to clean (same channel layout as the fitted data). - y : None - Ignored. - - Returns - ------- - X_clean : Raw | Epochs | ndarray - Cleaned data in the same format as the input. - """ - check_is_fitted(self, "spatial_filter_") - data, _sfreq, mne_type, orig_inst, picks, _names = extract_data_from_mne( - X, auto_pick=True + def transform(self, X: Any) -> Any: + """Apply the frozen spatial operator without refitting.""" + check_is_fitted( + self, + attributes=["spatial_filter_", "_n_channels_", "_fit_ch_names_"], ) - if mne_type == "epochs": - cleaned = np.empty_like(data, dtype=float) - for e in range(data.shape[0]): - cleaned[e] = self.spatial_filter_ @ np.asarray(data[e], dtype=float) - else: - cleaned = self.spatial_filter_ @ np.asarray(data, dtype=float) - return reconstruct_mne_object( - cleaned, orig_inst, mne_type, picks=picks, verbose=False + data, _, mne_type, orig_inst, picks, _ = extract_data_from_mne( + X, + ch_names=self._fit_ch_names_, + auto_pick=True, ) + data = np.asarray(data, dtype=float) + if not np.all(np.isfinite(data)): + raise ValueError("X must contain only finite values.") + n_channels = data.shape[1] if data.ndim == 3 else data.shape[0] + if n_channels != self._n_channels_: + raise ValueError( + "X has a different number of channels than the fitted data " + f"({n_channels} vs {self._n_channels_})." + ) + cleaned = _apply_spatial_filter(data, self.spatial_filter_) + return _reconstruct_like(cleaned, orig_inst, mne_type, picks) - def fit_transform(self, X: Any, y=None, mask: np.ndarray | None = None, **fit_params) -> Any: - """Fit on ``X`` and apply to ``X`` in one step. - - Parameters - ---------- - X : Raw | Epochs | ndarray - Input data. - y : None - Ignored. - mask : ndarray of bool | None - Optional artifact-segment mask (see :meth:`fit`). - **fit_params - Ignored. - - Returns - ------- - X_clean : Raw | Epochs | ndarray - Cleaned data. - """ - return self.fit(X, y, mask=mask).transform(X) - + def fit_transform( + self, + X: Any, + y=None, + *, + artifact_mask: np.ndarray | None = None, + clean_reference: Any | None = None, + **fit_params, + ) -> Any: + """Fit the operator and apply it to ``X``.""" + if fit_params: + unknown = ", ".join(sorted(fit_params)) + raise TypeError(f"Unexpected fit parameters: {unknown}.") + return self.fit( + X, + y, + artifact_mask=artifact_mask, + clean_reference=clean_reference, + ).transform(X) + + +# Short established acronym retained as a documented compatibility alias. +MWF = MultichannelWienerFilter + + +__all__ = [ + "MWF", + "MultichannelWienerFilter", + "compute_mwf", + "hf_power_mask", + "mwf_filter", +] diff --git a/tests/test_mwf.py b/tests/test_mwf.py index c2715b2f..de9ef346 100644 --- a/tests/test_mwf.py +++ b/tests/test_mwf.py @@ -1,218 +1,348 @@ -"""Tests for the mne_denoise.mwf module (multi-channel Wiener filter).""" +"""Tests for GEVD multi-channel Wiener filtering.""" from __future__ import annotations +import mne import numpy as np import pytest - -from mne_denoise.mwf import MWF, compute_mwf, hf_power_mask, mwf_filter - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - - -@pytest.fixture() -def rng(): - """Shared random generator.""" - return np.random.default_rng(0) - - -@pytest.fixture() -def burst_data(rng): - """Synthetic EEG with intermittent broadband high-frequency bursts. - - Returns ``(X, sfreq, burst)`` where ``burst`` is the boolean sample mask of - the artifact segments. The clean background is low-frequency; the artifact is - a broadband HF burst injected into a subset of samples. - """ +from sklearn.base import clone +from sklearn.exceptions import NotFittedError + +from mne_denoise.mwf import ( + MWF, + MultichannelWienerFilter, + compute_mwf, + hf_power_mask, + mwf_filter, +) + + +@pytest.fixture(scope="module") +def contaminated_data(): + """Known clean signal plus a marked spatial high-frequency artifact.""" + rng = np.random.default_rng(42) sfreq = 250.0 - n_times = 4000 - n_ch = 12 - t = np.arange(n_times) / sfreq - - # Low-frequency neural background (well separated from the HF artifact). - neural = np.vstack( - [np.sin(2 * np.pi * f * t + rng.uniform(0, 2 * np.pi)) for f in (6.0, 8.0, 10.0)] + n_channels, n_samples = 8, 4000 + times = np.arange(n_samples) / sfreq + sources = np.vstack( + [ + np.sin(2 * np.pi * 7.0 * times), + np.sin(2 * np.pi * 10.0 * times + 0.4), + np.sin(2 * np.pi * 13.0 * times + 0.8), + ] + ) + clean = rng.standard_normal((n_channels, 3)) @ sources + clean += 0.05 * rng.standard_normal(clean.shape) + artifact_mask = np.zeros(n_samples, dtype=bool) + for start in (400, 1400, 2400, 3300): + artifact_mask[start : start + 250] = True + artifact_topography = rng.standard_normal(n_channels) + artifact_topography /= np.linalg.norm(artifact_topography) + artifact_waveform = rng.standard_normal(n_samples) + artifact_waveform[~artifact_mask] = 0.0 + contaminated = clean + 8.0 * np.outer(artifact_topography, artifact_waveform) + return contaminated, clean, artifact_mask, artifact_topography, sfreq + + +def test_canonical_name_and_alias(): + """MWF is the documented alias of the canonical class.""" + assert MWF is MultichannelWienerFilter + assert isinstance(MWF(), MultichannelWienerFilter) + + +def test_estimator_is_cloneable(): + """Constructor parameters follow sklearn cloning semantics.""" + estimator = MultichannelWienerFilter(rank=3, reg=1e-4, treat_nan="clean") + cloned = clone(estimator) + assert cloned.get_params() == estimator.get_params() + + +def test_explicit_mask_is_required(contaminated_data): + """MWF does not silently invent an artifact definition.""" + data = contaminated_data[0] + with pytest.raises(ValueError, match="explicit artifact_mask"): + compute_mwf(data) + with pytest.raises(ValueError, match="explicit artifact_mask"): + MultichannelWienerFilter().fit(data) + + +def test_hf_power_mask_flags_bursts(contaminated_data): + """The opt-in HF heuristic enriches the known burst intervals.""" + data, _, true_mask, _, sfreq = contaminated_data + mask = hf_power_mask(data, sfreq, hf_hz=20.0, quantile=0.7) + assert mask.shape == (data.shape[1],) + assert mask.dtype == bool + assert (mask & true_mask).sum() / mask.sum() > 0.8 + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"sfreq": 0.0}, "sfreq"), + ({"sfreq": 250.0, "hf_hz": 125.0}, "Nyquist"), + ({"sfreq": 250.0, "quantile": 1.0}, "quantile"), + ({"sfreq": 250.0, "smooth_s": 0.0}, "smooth_s"), + ], +) +def test_hf_power_mask_validates_operating_point(kwargs, match): + """Physical mask parameters have explicit admissible ranges.""" + with pytest.raises(ValueError, match=match): + hf_power_mask(np.ones((3, 100)), **kwargs) + + +def test_hf_power_mask_rejects_too_short_data(): + """Short input fails clearly instead of leaking scipy's pad error.""" + with pytest.raises(ValueError, match="too short"): + hf_power_mask(np.ones((3, 8)), 250.0) + + +def test_compute_mwf_diagnostics_and_attenuation(contaminated_data): + """Known artifact energy falls and fit diagnostics remain explicit.""" + data, _, mask, topography, _ = contaminated_data + cleaned, diagnostics = compute_mwf(data, mask) + before = np.var(topography @ data[:, mask]) + after = np.var(topography @ cleaned[:, mask]) + assert after < 0.2 * before + assert diagnostics["artifact_samples"] == int(mask.sum()) + assert diagnostics["clean_samples"] == int((~mask).sum()) + assert diagnostics["rank_requested"] == "positive" + assert diagnostics["rank_used"] >= 1 + assert diagnostics["spatial_filter"].shape == (data.shape[0], data.shape[0]) + + +def test_mwf_filter_matches_compute_mwf(contaminated_data): + """The compact functional API delegates to the same implementation.""" + data, _, mask, _, _ = contaminated_data + expected, _ = compute_mwf(data, mask, rank=2) + observed = mwf_filter(data, mask, rank=2) + np.testing.assert_allclose(observed, expected) + + +def test_full_rank_matches_covariance_ratio(): + """Full-rank GEVD equals the regularized covariance-ratio MWF.""" + rng = np.random.default_rng(1) + data = rng.standard_normal((5, 3000)) + mask = np.zeros(data.shape[1], dtype=bool) + mask[:1200] = True + data[:, mask] += np.outer(rng.standard_normal(5), rng.standard_normal(mask.sum())) + _, diagnostics = compute_mwf(data, mask, rank="full", reg=1e-4) + artifact_cov = np.cov(data[:, mask]) + clean_cov = np.cov(data[:, ~mask]) + scale = max(np.trace(artifact_cov) / 5, np.trace(clean_cov) / 5) + artifact_cov += 1e-4 * scale * np.eye(5) + clean_cov += 1e-4 * scale * np.eye(5) + expected = clean_cov @ np.linalg.inv(artifact_cov) + np.testing.assert_allclose( + diagnostics["spatial_filter"], expected, rtol=1e-8, atol=1e-10 ) - M_neural = rng.standard_normal((n_ch, 3)) - X = M_neural @ neural - - # Broadband HF bursts on ~30% of samples. - burst = np.zeros(n_times, dtype=bool) - for start in (400, 1500, 2600, 3300): - burst[start : start + 300] = True - hf = rng.standard_normal((n_ch, n_times)) * 3.0 - from scipy.signal import butter, filtfilt - - b, a = butter(4, 40.0 / (0.5 * sfreq), btype="high") - hf = filtfilt(b, a, hf, axis=-1) - X = X + hf * burst - return X, sfreq, burst - - -def _band_power(X, sfreq, fmin, fmax): - spec = np.abs(np.fft.rfft(np.atleast_2d(X), axis=-1)) ** 2 - freqs = np.fft.rfftfreq(X.shape[-1], 1.0 / sfreq) - band = (freqs >= fmin) & (freqs < fmax) - return float(spec[:, band].sum()) - - -# --------------------------------------------------------------------------- -# Functional API -# --------------------------------------------------------------------------- - - -def test_hf_power_mask_flags_bursts(burst_data): - """The HF-power mask overlaps the injected burst samples.""" - X, sfreq, burst = burst_data - mask = hf_power_mask(X, sfreq, hf_hz=20.0, quantile=0.6) - assert mask.shape == (X.shape[-1],) - # Detected artifact samples are enriched inside the true bursts. - overlap = (mask & burst).sum() / max(mask.sum(), 1) - assert overlap > 0.5 - - -def test_compute_mwf_shapes_and_info(burst_data): - """compute_mwf returns cleaned data + a diagnostics dict.""" - X, sfreq, _burst = burst_data - cleaned, info = compute_mwf(X, sfreq) - assert cleaned.shape == X.shape - assert info["mask"].shape == (X.shape[-1],) - assert 0.0 <= info["artifact_fraction"] <= 1.0 - - -def test_compute_mwf_rejects_1d(): - """A 1-D input raises a clear error.""" - with pytest.raises(ValueError, match="2-D"): - compute_mwf(np.zeros(100), 250.0) - - -def test_mwf_filter_identity_when_covariances_match(rng): - """With a random mask on stationary data the filter is ~identity (no over-clean).""" - X = rng.standard_normal((8, 4000)) - mask = np.zeros(4000, dtype=bool) - mask[::2] = True # arbitrary split of stationary data - cleaned = mwf_filter(X, mask, reg=1e-6) - # Stationary data -> R_clean ~ R_artifact -> filter ~ identity. - assert np.corrcoef(cleaned.ravel(), X.ravel())[0, 1] > 0.9 - - -# --------------------------------------------------------------------------- -# MWF estimator -# --------------------------------------------------------------------------- - - -def test_mwf_fit_transform_numpy_shape(burst_data): - """fit_transform on a NumPy array returns an array of the same shape.""" - X, sfreq, _burst = burst_data - cleaned = MWF(sfreq=sfreq).fit_transform(X) - assert isinstance(cleaned, np.ndarray) - assert cleaned.shape == X.shape - - -def test_mwf_reduces_burst_hf_power(burst_data): - """MWF attenuates the broadband HF burst power.""" - X, sfreq, _burst = burst_data - cleaned = MWF(sfreq=sfreq, hf_hz=20.0, quantile=0.6).fit_transform(X) - hf_before = _band_power(X, sfreq, 40.0, 120.0) - hf_after = _band_power(cleaned, sfreq, 40.0, 120.0) - assert hf_after < hf_before - - -def test_mwf_fitted_attributes(burst_data): - """Fitted attributes are populated with correct shapes.""" - X, sfreq, _burst = burst_data - est = MWF(sfreq=sfreq).fit(X) - assert est.spatial_filter_.shape == (X.shape[0], X.shape[0]) - assert est.artifact_mask_.shape == (X.shape[-1],) - assert 0.0 <= est.artifact_fraction_ <= 1.0 - - -def test_mwf_leakage_split_applies_operator(rng): - """transform applies the operator learned in fit (train != eval).""" - train = rng.standard_normal((8, 4000)) - evalu = rng.standard_normal((8, 1000)) - mask = np.zeros(4000, dtype=bool) - mask[:1500] = True - est = MWF().fit(train, mask=mask) # explicit mask -> no sfreq needed - cleaned = est.transform(evalu) - np.testing.assert_allclose(cleaned, est.spatial_filter_ @ evalu) - - -def test_mwf_requires_sfreq_for_array_without_mask(burst_data): - """Array input without sfreq or mask raises a clear error.""" - X, _sfreq, _burst = burst_data - with pytest.raises(ValueError, match="sfreq is required"): - MWF().fit(X) - -def test_mwf_transform_before_fit_raises(rng): - """transform before fit raises NotFittedError.""" - from sklearn.exceptions import NotFittedError +def test_channel_means_are_preserved(contaminated_data): + """MWF follows the reference mean-subtract/apply/restore convention.""" + data, _, mask, _, _ = contaminated_data + shifted = data + np.arange(data.shape[0])[:, None] * 10.0 + cleaned, _ = compute_mwf(shifted, mask) + np.testing.assert_allclose(cleaned.mean(axis=1), shifted.mean(axis=1), atol=1e-12) + + +def test_global_unit_rescaling_preserves_operator(contaminated_data): + """Relative regularization is invariant to a shared physical-unit scale.""" + data, _, mask, _, _ = contaminated_data + cleaned, diagnostics = compute_mwf(data, mask) + scaled, scaled_diagnostics = compute_mwf(data * 1e6, mask) + np.testing.assert_allclose( + diagnostics["spatial_filter"], + scaled_diagnostics["spatial_filter"], + rtol=1e-8, + atol=1e-10, + ) + np.testing.assert_allclose(scaled, cleaned * 1e6, rtol=1e-8, atol=1e-6) + + +def test_rank_deficiency_is_regularized(): + """Default loading supports rank-deficient data while zero loading is explicit.""" + rng = np.random.default_rng(2) + latent = rng.standard_normal((2, 1000)) + mixing = rng.standard_normal((6, 2)) + data = mixing @ latent + mask = np.zeros(1000, dtype=bool) + mask[:400] = True + cleaned, diagnostics = compute_mwf(data, mask, reg=1e-6) + assert np.all(np.isfinite(cleaned)) + assert diagnostics["clean_covariance_rank"] < data.shape[0] + with pytest.raises(ValueError, match="positive definite"): + compute_mwf(data, mask, reg=0.0) + + +@pytest.mark.parametrize("bad_mask", [np.zeros(10), np.full(4000, 2), ["x"] * 4000]) +def test_artifact_mask_validation(contaminated_data, bad_mask): + """Mask shape and values cannot silently change the training regime.""" + data = contaminated_data[0] + with pytest.raises((TypeError, ValueError), match="artifact_mask"): + compute_mwf(data, bad_mask) + + +def test_nan_mask_policies(contaminated_data): + """Ignored mask samples stay out of both covariance estimates.""" + data, _, mask, _, _ = contaminated_data + ternary = mask.astype(float) + ternary[-200:] = np.nan + _, ignored = compute_mwf(data, ternary, treat_nan="ignore") + _, clean = compute_mwf(data, ternary, treat_nan="clean") + assert ignored["ignored_samples"] == 200 + assert clean["ignored_samples"] == 0 + assert clean["clean_samples"] == ignored["clean_samples"] + 200 + + +def test_clean_reference_contract(contaminated_data): + """A clean reference can supply the clean covariance explicitly.""" + data, clean, _, _, _ = contaminated_data + cleaned, diagnostics = compute_mwf(data, clean_reference=clean) + assert cleaned.shape == data.shape + assert diagnostics["used_clean_reference"] is True + assert diagnostics["artifact_samples"] == data.shape[1] + + +def test_inadmissible_inputs_fail_clearly(contaminated_data): + """Non-finite, single-channel, and insufficient masks are rejected.""" + data, _, mask, _, _ = contaminated_data + with pytest.raises(ValueError, match="at least 2 channels"): + compute_mwf(data[:1], mask) + nonfinite = data.copy() + nonfinite[0, 0] = np.nan + with pytest.raises(ValueError, match="finite"): + compute_mwf(nonfinite, mask) + with pytest.raises(ValueError, match="two artifact"): + compute_mwf(data, np.zeros(data.shape[1], dtype=bool)) + + +def test_estimator_fit_transform_and_frozen_operator(contaminated_data): + """Transform applies the fitted operator without evaluating a new mask.""" + data, _, mask, _, _ = contaminated_data + estimator = MultichannelWienerFilter(rank=2).fit(data, artifact_mask=mask) + evaluation = data[:, :500] * 0.5 + observed = estimator.transform(evaluation) + means = evaluation.mean(axis=1, keepdims=True) + expected = estimator.spatial_filter_ @ (evaluation - means) + means + np.testing.assert_allclose(observed, expected) + assert estimator.artifact_mask_.shape == mask.shape + assert estimator.selected_components_.size == 2 + + +def test_fit_transform_accepts_mask(contaminated_data): + """fit_transform forwards explicit fit assets without leakage ambiguity.""" + data, _, mask, _, _ = contaminated_data + cleaned = MultichannelWienerFilter().fit_transform(data, artifact_mask=mask) + assert cleaned.shape == data.shape + + +def test_transform_before_fit_and_channel_mismatch(contaminated_data): + """Estimator state and sensor dimensionality are enforced.""" + data, _, mask, _, _ = contaminated_data with pytest.raises(NotFittedError): - MWF().transform(rng.standard_normal((8, 100))) - - -# --------------------------------------------------------------------------- -# MNE round-trip -# --------------------------------------------------------------------------- - - -def test_mwf_mne_raw_roundtrip_infers_sfreq(burst_data): - """fit_transform on an MNE Raw returns a Raw of identical shape; sfreq inferred.""" - mne = pytest.importorskip("mne") - X, sfreq, _burst = burst_data - info = mne.create_info([f"EEG{i:02d}" for i in range(X.shape[0])], sfreq, "eeg") - raw = mne.io.RawArray(X, info, verbose=False) - - cleaned = MWF(hf_hz=20.0, quantile=0.6).fit_transform(raw) # sfreq from info - assert isinstance(cleaned, mne.io.BaseRaw) - assert cleaned.get_data().shape == X.shape - assert not np.allclose(cleaned.get_data(), X) - - -# --------------------------------------------------------------------------- -# Edge cases -# --------------------------------------------------------------------------- - - -def test_hf_power_mask_short_signal_length(rng): - """The mask length equals n_times even when the smoothing window is longer.""" - # 20 samples at 250 Hz -> default smooth window (25 samples) exceeds length. - X = rng.standard_normal((8, 20)) - mask = hf_power_mask(X, 250.0) - assert mask.shape == (20,) - - -def test_mwf_short_signal_no_crash(rng): - """compute_mwf / MWF.fit_transform do not crash on very short signals.""" - X = rng.standard_normal((8, 20)) - cleaned, info = compute_mwf(X, 250.0) - assert cleaned.shape == X.shape - assert info["mask"].shape == (20,) - cleaned2 = MWF(sfreq=250.0).fit_transform(X) - assert cleaned2.shape == X.shape - - -def test_mwf_single_channel_is_identity(rng): - """A single channel is degenerate for a spatial filter -> returned unchanged.""" - X = rng.standard_normal((1, 2000)) - cleaned = MWF(sfreq=250.0).fit_transform(X) - assert cleaned.shape == X.shape - np.testing.assert_allclose(cleaned, X) - # functional API too - cleaned_f, _ = compute_mwf(X, 250.0) - assert cleaned_f.shape == X.shape - - -def test_mwf_int_mask_matches_bool_mask(rng): - """An integer 0/1 mask gives the same result as the equivalent bool mask.""" - X = rng.standard_normal((8, 2000)) - bool_mask = np.zeros(2000, dtype=bool) - bool_mask[:800] = True - int_mask = bool_mask.astype(int) - cleaned_bool, _ = compute_mwf(X, 250.0, mask=bool_mask) - cleaned_int, _ = compute_mwf(X, 250.0, mask=int_mask) - np.testing.assert_allclose(cleaned_bool, cleaned_int) + MultichannelWienerFilter().transform(data) + estimator = MultichannelWienerFilter().fit(data, artifact_mask=mask) + with pytest.raises(ValueError, match="different number of channels"): + estimator.transform(data[:-1]) + + +def test_explicit_hf_strategy_requires_consistent_sfreq(contaminated_data): + """Automatic mask creation is opt-in and its sampling rate is checked.""" + data, _, _, _, sfreq = contaminated_data + estimator = MultichannelWienerFilter( + mask_strategy="hf_power", sfreq=sfreq, quantile=0.7 + ).fit(data) + assert estimator.fit_diagnostics_["mask_strategy"] == "hf_power" + with pytest.raises(ValueError, match="sfreq is required"): + MultichannelWienerFilter(mask_strategy="hf_power").fit(data) + + +def _make_raw(contaminated_data): + """Create Raw with EEG plus an untouched stimulus channel and metadata.""" + data, _, mask, _, sfreq = contaminated_data + eeg = data[:4] + stim = np.zeros((1, data.shape[1])) + stim[0, 123] = 1 + names = [f"EEG{index:02d}" for index in range(4)] + ["STI 014"] + info = mne.create_info(names, sfreq, ["eeg"] * 4 + ["stim"]) + info["bads"] = ["EEG01"] + raw = mne.io.RawArray(np.vstack([eeg, stim]), info, first_samp=100, verbose=False) + raw.set_annotations(mne.Annotations([1.0], [0.2], ["test"]), emit_warning=False) + return raw, mask + + +def test_raw_roundtrip_preserves_metadata_and_unpicked_channels(contaminated_data): + """Raw subtype, first sample, annotations, bads, and stim data survive.""" + raw, mask = _make_raw(contaminated_data) + estimator = MultichannelWienerFilter().fit(raw, artifact_mask=mask) + cleaned = estimator.transform(raw) + assert type(cleaned) is type(raw) + assert cleaned.first_samp == raw.first_samp + assert cleaned.info["bads"] == raw.info["bads"] + assert cleaned.annotations == raw.annotations + np.testing.assert_array_equal( + cleaned.get_data(picks=["STI 014"]), raw.get_data(picks=["STI 014"]) + ) + assert not np.allclose(cleaned.get_data(picks="eeg"), raw.get_data(picks="eeg")) + + +def test_mne_channel_alignment_and_sfreq_mismatch(contaminated_data): + """Named channels are aligned and conflicting physical time units fail.""" + raw, mask = _make_raw(contaminated_data) + estimator = MultichannelWienerFilter().fit(raw, artifact_mask=mask) + reordered = raw.copy().reorder_channels(list(reversed(raw.ch_names))) + cleaned = estimator.transform(reordered) + assert cleaned.ch_names == reordered.ch_names + with pytest.raises(ValueError, match="does not match MNE metadata"): + MultichannelWienerFilter(sfreq=500.0).fit(raw, artifact_mask=mask) + + +def test_epochs_roundtrip_and_epoch_mask(contaminated_data): + """Epoch masks flatten consistently while events and metadata survive.""" + pandas = pytest.importorskip("pandas") + data, _, mask, _, sfreq = contaminated_data + epochs_data = data[:4, :2000].reshape(4, 5, 400).transpose(1, 0, 2) + epoch_mask = mask[:2000].reshape(5, 400) + info = mne.create_info([f"EEG{i}" for i in range(4)], sfreq, "eeg") + events = np.column_stack([np.arange(5) * 500, np.zeros(5, int), np.ones(5, int)]) + metadata = pandas.DataFrame({"trial": np.arange(5)}) + epochs = mne.EpochsArray( + epochs_data, + info, + events=events, + event_id={"event": 1}, + metadata=metadata, + verbose=False, + ) + cleaned = MultichannelWienerFilter().fit_transform(epochs, artifact_mask=epoch_mask) + assert type(cleaned) is type(epochs) + np.testing.assert_array_equal(cleaned.events, epochs.events) + assert cleaned.event_id == epochs.event_id + assert cleaned.metadata.equals(epochs.metadata) + assert cleaned.get_data().shape == epochs.get_data().shape + + +def test_evoked_roundtrip(contaminated_data): + """Evoked comment, nave, time origin, and data shape survive.""" + data, _, mask, _, sfreq = contaminated_data + info = mne.create_info([f"EEG{i}" for i in range(4)], sfreq, "eeg") + evoked = mne.EvokedArray( + data[:4], info, tmin=-0.2, nave=12, comment="condition", verbose=False + ) + cleaned = MultichannelWienerFilter().fit_transform(evoked, artifact_mask=mask) + assert type(cleaned) is type(evoked) + assert cleaned.comment == evoked.comment + assert cleaned.nave == evoked.nave + assert cleaned.first == evoked.first + assert cleaned.data.shape == evoked.data.shape + + +def test_mne_clean_reference_alignment_and_sfreq(contaminated_data): + """Reference channels align by name and reference sfreq must match.""" + raw, _ = _make_raw(contaminated_data) + reference = raw.copy().reorder_channels(list(reversed(raw.ch_names))) + estimator = MultichannelWienerFilter().fit(raw, clean_reference=reference) + assert estimator.fit_diagnostics_["used_clean_reference"] is True + mismatch = reference.copy().resample(125.0, verbose=False) + with pytest.raises(ValueError, match="sampling frequency must match"): + MultichannelWienerFilter().fit(raw, clean_reference=mismatch) From 380c229fa195a326bf72bffc338e25d0f08a3b8b Mon Sep 17 00:00:00 2001 From: Sina Esmaeili Date: Sat, 1 Aug 2026 04:42:47 -0400 Subject: [PATCH 3/4] test(mwf): allow cross-BLAS eigensolver tolerance --- tests/test_mwf.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_mwf.py b/tests/test_mwf.py index de9ef346..ef6369a1 100644 --- a/tests/test_mwf.py +++ b/tests/test_mwf.py @@ -156,7 +156,9 @@ def test_global_unit_rescaling_preserves_operator(contaminated_data): rtol=1e-8, atol=1e-10, ) - np.testing.assert_allclose(scaled, cleaned * 1e6, rtol=1e-8, atol=1e-6) + # Generalized eigensolvers vary by a few ulps across BLAS implementations; + # this still bounds the unit-rescaling error below one part in 10 million. + np.testing.assert_allclose(scaled, cleaned * 1e6, rtol=1e-7, atol=2e-6) def test_rank_deficiency_is_regularized(): From c05db1ac04b942b23db9f28f1e1e1693b5e0ffba Mon Sep 17 00:00:00 2001 From: Sina Esmaeili Date: Sat, 1 Aug 2026 05:39:58 -0400 Subject: [PATCH 4/4] test(mwf): cover execution preconditions --- tests/test_mwf.py | 86 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/tests/test_mwf.py b/tests/test_mwf.py index ef6369a1..41c8ea7c 100644 --- a/tests/test_mwf.py +++ b/tests/test_mwf.py @@ -2,6 +2,8 @@ from __future__ import annotations +import logging + import mne import numpy as np import pytest @@ -15,6 +17,7 @@ hf_power_mask, mwf_filter, ) +from mne_denoise.mwf.core import _apply_spatial_filter, _compute_operator @pytest.fixture(scope="module") @@ -348,3 +351,86 @@ def test_mne_clean_reference_alignment_and_sfreq(contaminated_data): mismatch = reference.copy().resample(125.0, verbose=False) with pytest.raises(ValueError, match="sampling frequency must match"): MultichannelWienerFilter().fit(raw, clean_reference=mismatch) + + +def test_input_and_operating_point_validation_branches(contaminated_data): + """Invalid arrays, masks, ranks, and explicit strategies fail at the boundary.""" + data, _, mask, _, sfreq = contaminated_data + + with pytest.raises(TypeError, match="numeric array"): + compute_mwf("not numeric") + with pytest.raises(ValueError, match="shape"): + compute_mwf(np.ones(10), np.ones(10)) + with pytest.raises(ValueError, match="at least 2 samples"): + compute_mwf(np.ones((2, 1)), np.ones(1)) + with pytest.raises(ValueError, match="finite real"): + hf_power_mask(data, True) + with pytest.raises(ValueError, match="treat_nan"): + compute_mwf(data, mask, treat_nan="invalid") + + infinite_mask = mask.astype(float) + infinite_mask[0] = np.inf + with pytest.raises(ValueError, match="0, 1, or NaN"): + compute_mwf(data, infinite_mask) + + for bad_rank, error in ( + ("invalid", ValueError), + (True, TypeError), + (100, ValueError), + ): + with pytest.raises(error, match="rank"): + compute_mwf(data, mask, rank=bad_rank) + with pytest.raises(ValueError, match="reg must be >="): + compute_mwf(data, mask, reg=-1.0) + with pytest.raises(ValueError, match="Pass artifact_mask or mask_strategy"): + compute_mwf(data, mask, mask_strategy="hf_power", sfreq=sfreq) + with pytest.raises(ValueError, match="mask_strategy"): + compute_mwf(data, mask_strategy="invalid") + + +def test_mask_policy_and_training_preconditions(contaminated_data): + """Mask conversion and covariance preconditions remain explicit.""" + data, _, mask, _, _ = contaminated_data + ternary = mask.astype(float) + ternary[-200:] = np.nan + _, diagnostics = compute_mwf(data, ternary, treat_nan="artifact") + assert diagnostics["artifact_samples"] == int(mask.sum()) + 200 + assert diagnostics["ignored_samples"] == 0 + + with pytest.raises(ValueError, match="same number of channels"): + compute_mwf(data, mask, clean_reference=data[:-1]) + almost_all_artifact = np.ones(data.shape[1], dtype=bool) + almost_all_artifact[0] = False + with pytest.raises(ValueError, match="two clean samples"): + compute_mwf(data, almost_all_artifact) + + zero_data = np.zeros((2, 6)) + zero_mask = np.array([1, 1, 1, 0, 0, 0], dtype=bool) + with pytest.raises(ValueError, match="zero numerical energy"): + compute_mwf(zero_data, zero_mask) + + +def test_internal_operator_shape_preconditions(): + """The mathematical core rejects mismatched or undersampled training arrays.""" + kwargs = {"rank": "positive", "artifact_weight": 1.0, "reg": 1e-6} + with pytest.raises(ValueError, match="same channel count"): + _compute_operator(np.ones((2, 4)), np.ones((3, 4)), **kwargs) + with pytest.raises(ValueError, match="at least two"): + _compute_operator(np.ones((2, 1)), np.ones((2, 4)), **kwargs) + with pytest.raises(ValueError, match="2D or 3D"): + _apply_spatial_filter(np.ones(4), np.eye(2)) + + +def test_estimator_diagnostic_and_error_paths(contaminated_data, caplog): + """Logging, frozen evaluation checks, and unknown fit assets are covered.""" + data, _, mask, _, _ = contaminated_data + with caplog.at_level(logging.INFO, logger="mne_denoise.mwf.core"): + estimator = MultichannelWienerFilter(verbose=True).fit(data, artifact_mask=mask) + assert "MWF fit:" in caplog.text + + nonfinite = data[:, :20].copy() + nonfinite[0, 0] = np.nan + with pytest.raises(ValueError, match="finite"): + estimator.transform(nonfinite) + with pytest.raises(TypeError, match="Unexpected fit parameters: unknown"): + estimator.fit_transform(data, artifact_mask=mask, unknown=True)