diff --git a/docs/api.rst b/docs/api.rst index e00aeaf5..426d6ccf 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -63,6 +63,20 @@ iCanClean mne_denoise.icanclean.ICanClean mne_denoise.icanclean.compute_icanclean +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/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..ba4f793d --- /dev/null +++ b/docs/mwf.md @@ -0,0 +1,108 @@ +# Multi-channel Wiener filtering + +`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. + +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. + +## 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) +``` + +`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. + +## Independent clean reference + +A separate clean recording can supply the clean covariance. If no mask is +provided, all samples in `X` train the artifact-present covariance: + +```python +mwf.fit(artifact_training_raw, clean_reference=clean_reference_raw) +``` + +The training and reference data must have the same channels, channel scaling, +physical units, and—when both are MNE objects—the same sampling frequency. + +## Optional high-frequency mask authoring + +`hf_power_mask()` is a convenience heuristic, not part of the reference MWF +algorithm. It can be used explicitly: + +```python +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) +``` + +Or requested as an explicit estimator strategy: + +```python +mwf = MultichannelWienerFilter( + mask_strategy="hf_power", + sfreq=250.0, + hf_hz=20.0, + quantile=0.7, +) +mwf.fit(train_data) +``` + +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. + +## GEVD rank and diagnostics + +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. + +## Evidence boundary + +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. Authors' MATLAB implementation: + https://github.com/exporl/mwf-artifact-removal 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/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 new file mode 100644 index 00000000..3e1e1b2e --- /dev/null +++ b/mne_denoise/mwf/__init__.py @@ -0,0 +1,23 @@ +"""Semi-supervised multi-channel Wiener filtering. + +``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, + 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 new file mode 100644 index 00000000..d7ff0701 --- /dev/null +++ b/mne_denoise/mwf/core.py @@ -0,0 +1,790 @@ +"""GEVD multi-channel Wiener filtering for semi-supervised EEG denoising. + +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. + +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 +---------- +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 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 + +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: + """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_samples) + Multichannel signal in any consistent physical unit. + sfreq : float + Sampling frequency in Hz. + hf_hz : float + High-pass cutoff in Hz. It must lie strictly below Nyquist. + quantile : float + Fractional power quantile strictly between zero and one. Samples above + this threshold are marked as artifact. + smooth_s : float + Positive moving-average duration in seconds. + + Returns + ------- + artifact_mask : ndarray of bool, shape (n_samples,) + True for samples selected as artifact-present. + + Notes + ----- + The output depends on ``sfreq``, ``hf_hz``, ``quantile``, and ``smooth_s``; + these values are part of the MWF operating point, not universal defaults. + """ + 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, + 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, + treat_nan: Literal["ignore", "artifact", "clean"] = "ignore", +) -> tuple[np.ndarray, dict[str, Any]]: + """Fit and apply a zero-delay GEVD multi-channel Wiener filter. + + Parameters + ---------- + 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 + Non-negative relative diagonal loading applied to both covariances. + treat_nan : {'ignore', 'artifact', 'clean'} + Interpretation of NaNs in ``artifact_mask``. + + Returns + ------- + 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. + """ + 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 + ---------- + 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 + 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 + Enable a concise fit summary when truthy. + + Attributes + ---------- + spatial_filter_ : ndarray, shape (n_channels, n_channels) + 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, + 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.smooth_s = smooth_s + self.verbose = verbose + + def _resolve_sfreq(self, sfreq_data: float | None) -> float | None: + """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, + ) + data = _as_2d_finite(data, name="X", min_channels=2) + return data, sfreq, extracted_names + + 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 | Evoked | ndarray + Training data. Epochs are concatenated in epoch-major time order. + y : None + Ignored. + 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 : MultichannelWienerFilter + Fitted estimator. + """ + 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( + "clean_reference sampling frequency must match the training data." + ) + + 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: %d artifact samples, %d clean samples, rank %d.", + diagnostics["artifact_samples"], + diagnostics["clean_samples"], + diagnostics["rank_used"], + ) + return self + + 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_"], + ) + 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, + *, + 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 new file mode 100644 index 00000000..41c8ea7c --- /dev/null +++ b/tests/test_mwf.py @@ -0,0 +1,436 @@ +"""Tests for GEVD multi-channel Wiener filtering.""" + +from __future__ import annotations + +import logging + +import mne +import numpy as np +import pytest +from sklearn.base import clone +from sklearn.exceptions import NotFittedError + +from mne_denoise.mwf import ( + MWF, + MultichannelWienerFilter, + compute_mwf, + hf_power_mask, + mwf_filter, +) +from mne_denoise.mwf.core import _apply_spatial_filter, _compute_operator + + +@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_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 + ) + + +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, + ) + # 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(): + """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): + 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) + + +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)