Source code for mne_denoise.sns.core

"""Sensor Noise Suppression (SNS) for channel-specific noise removal.

SNS [1]_ suppresses noise that is specific to an
individual sensor by regenerating every channel from a least-squares projection
onto correlated neighbouring channels.  The signal of interest must therefore
be spatially redundant, while the targeted noise must be sensor-specific.

The module provides a covariance-level primitive, a one-shot array function,
and a fitted estimator.  The estimator learns both its centering statistics and
spatial operator on the training set, so transforming a sample does not depend
on the other samples supplied in the same call.

References
----------
.. [1] de Cheveigné, A., & Simon, J. Z. (2008). Sensor noise suppression.
       Journal of Neuroscience Methods, 168(1), 195-202.
       https://doi.org/10.1016/j.jneumeth.2007.09.012
"""

from __future__ import annotations

import logging
from numbers import Integral, Real
from typing import Any

import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.utils.validation import check_is_fitted

from .._covariance import compute_covariance
from .._logging import set_log_level_from_verbose
from .._spatial import (
    apply_spatial_transform,
    continuous_to_epochs,
    epochs_to_continuous,
)
from .._validation import (
    check_channel_first_data,
    check_channel_layout,
    check_chunk_size,
)
from ..utils import extract_data_from_mne, reconstruct_mne_object

logger = logging.getLogger(__name__)

_DEFAULT_RCOND = 1e-12


def _automatic_sample_mask(
    data: np.ndarray,
    manual_weight: np.ndarray,
    threshold: float | None,
) -> np.ndarray:
    """Reject samples with a large robust deviation in any channel."""
    if threshold is None:
        return np.ones(data.shape[1], dtype=np.float64)
    included = manual_weight > 0
    reference = data[:, included]
    center = np.median(reference, axis=1, keepdims=True)
    mad = np.median(np.abs(reference - center), axis=1, keepdims=True)
    scale = 1.4826 * mad
    fallback = np.std(reference, axis=1, keepdims=True)
    scale = np.where(scale > 0, scale, fallback)
    scale = np.where(scale > 0, scale, 1.0)
    max_abs_z = np.max(np.abs((data - center) / scale), axis=0)
    return (max_abs_z <= threshold).astype(np.float64)


[docs] def compute_sns_weights( cov: np.ndarray, n_neighbors: int = 0, skip: int = 0, *, rcond: float = _DEFAULT_RCOND, ) -> tuple[np.ndarray, int, np.ndarray]: """Compute the SNS spatial operator from a channel covariance matrix. Each channel is regenerated by a least-squares projection onto its most correlated neighbour channels. The channel itself is always excluded, so the diagonal of the resulting operator is zero [1]_. Parameters ---------- cov : ndarray, shape (n_channels, n_channels) Finite, symmetric, positive-semidefinite channel covariance matrix. n_neighbors : int, default=0 Number of neighbours used to regenerate each channel. Zero uses all available channels after applying ``skip``. skip : int, default=0 Number of the most-correlated neighbours to omit. This can be useful when adjacent sensors may share local noise. rcond : float, default=1e-12 Relative cutoff for the pseudoinverse of each neighbour covariance. Returns ------- weights : ndarray, shape (n_channels, n_channels) Spatial operator to apply to centered channel-first data. n_neighbors_used : int Effective number of neighbours after capping to those available. neighbor_ranks : ndarray, shape (n_channels,) Numerical rank of each selected neighbour covariance. Raises ------ TypeError If an integer parameter or ``rcond`` has an invalid type. ValueError If ``cov`` or an operating parameter is invalid. Notes ----- For centered data ``X``, the regenerated signal is ``weights @ X``. References ---------- .. [1] de Cheveigné, A., & Simon, J. Z. (2008). Sensor noise suppression. Journal of Neuroscience Methods, 168(1), 195-202. https://doi.org/10.1016/j.jneumeth.2007.09.012 """ cov = np.asarray(cov, dtype=np.float64) if cov.ndim != 2 or cov.shape[0] != cov.shape[1]: raise ValueError( "cov must be a square (n_channels, n_channels) matrix, " f"got shape {cov.shape}" ) if cov.shape[0] < 2: raise ValueError("SNS requires at least two channels") if not np.isfinite(cov).all(): raise ValueError("cov must contain only finite values") scale = max(float(np.max(np.abs(cov))), np.finfo(np.float64).tiny) tolerance = 100 * np.finfo(np.float64).eps * scale * cov.shape[0] if float(np.max(np.abs(cov - cov.T))) > tolerance: raise ValueError("cov must be symmetric") cov = (cov + cov.T) / 2.0 eigenvalues = np.linalg.eigvalsh(cov) eigen_scale = max(float(np.max(np.abs(eigenvalues))), np.finfo(float).tiny) eigen_tolerance = 100 * np.finfo(float).eps * eigen_scale * cov.shape[0] if float(eigenvalues.min()) < -eigen_tolerance: raise ValueError("cov must be positive semidefinite") for value, name in ((n_neighbors, "n_neighbors"), (skip, "skip")): if isinstance(value, bool) or not isinstance(value, Integral): raise TypeError(f"{name} must be a non-negative integer") if value < 0: raise ValueError(f"{name} must be a non-negative integer") if isinstance(rcond, bool) or not isinstance(rcond, Real): raise TypeError("rcond must be a finite number") rcond = float(rcond) if not np.isfinite(rcond) or not 0.0 < rcond < 1.0: raise ValueError("rcond must be finite and strictly between 0 and 1") n_channels = cov.shape[0] if skip > n_channels - 2: raise ValueError("skip must leave at least one candidate neighbor") max_neighbors = n_channels - int(skip) - 1 k_neighbors = ( max_neighbors if n_neighbors == 0 else min(int(n_neighbors), max_neighbors) ) standard_deviation = np.sqrt(np.clip(np.diag(cov), 0.0, None)) denominator = np.outer(standard_deviation, standard_deviation) denominator[denominator == 0.0] = 1.0 correlation = cov / denominator weights = np.zeros_like(cov) neighbor_ranks = np.zeros(n_channels, dtype=int) for channel in range(n_channels): order = np.argsort(correlation[:, channel] ** 2, kind="stable")[::-1] order = order[order != channel] neighbors = order[int(skip) : int(skip) + k_neighbors] neighbor_cov = cov[np.ix_(neighbors, neighbors)] singular_values = np.linalg.eigvalsh(neighbor_cov) cutoff = rcond * max(float(singular_values.max()), 0.0) neighbor_ranks[channel] = np.count_nonzero(singular_values > cutoff) weights[channel, neighbors] = ( np.linalg.pinv(neighbor_cov, rcond=rcond, hermitian=True)
[docs] @ cov[neighbors, channel] ) return weights, k_neighbors, neighbor_ranks
def compute_sns( X: np.ndarray, n_neighbors: int = 0, skip: int = 0, *, rcond: float = _DEFAULT_RCOND, preserve_mean: bool = False, n_iter: int = 1, outlier_threshold: float | None = None, chunk_size: int | None = None, sample_weight: np.ndarray | None = None, ) -> tuple[np.ndarray, dict[str, Any]]: """Learn and apply Sensor Noise Suppression to a channel-first array. This convenience function applies the algorithm described in [1]_ to the same data used to estimate its spatial operator. Use :class:`SNS` when fitting and transforming separate data. Parameters ---------- X : ndarray, shape (n_channels, n_times) | (n_epochs, n_channels, n_times) Multichannel continuous or epoched data. n_neighbors : int, default=0 Number of neighbours used to regenerate each channel. Zero uses all available channels after applying ``skip``. skip : int, default=0 Number of the most-correlated neighbours to omit. rcond : float, default=1e-12 Relative cutoff for local covariance pseudoinverses. preserve_mean : bool, default=False If True, add the fitted channel means back after regeneration. n_iter : int, default=1 Number of successive SNS projections to learn and compose. outlier_threshold : float | None, default=None Maximum robust channel-wise z-score allowed when learning the operator. Rejected samples are still transformed. None disables rejection. chunk_size : int | None, default=None Number of samples processed at once during covariance accumulation and operator application. None processes all samples together. sample_weight : ndarray, shape (n_times,) | (n_epochs, n_times) | None Non-negative fitting weight for each sample. Zero excludes a sample when learning the mean and operator, but not when applying the operator. Returns ------- X_clean : ndarray Sensor-noise-suppressed data with the same shape as ``X``. info : dict Fitted operator and numerical diagnostics. References ---------- .. [1] de Cheveigné, A., & Simon, J. Z. (2008). Sensor noise suppression. Journal of Neuroscience Methods, 168(1), 195-202. https://doi.org/10.1016/j.jneumeth.2007.09.012 """ X = check_channel_first_data(X, name="SNS") if not isinstance(preserve_mean, bool): raise TypeError("preserve_mean must be a bool") if isinstance(n_iter, bool) or not isinstance(n_iter, Integral): raise TypeError("n_iter must be a positive integer") if n_iter < 1: raise ValueError("n_iter must be a positive integer") chunk_size = check_chunk_size(chunk_size) if outlier_threshold is not None: if isinstance(outlier_threshold, bool) or not isinstance( outlier_threshold, Real ): raise TypeError("outlier_threshold must be a positive number or None") outlier_threshold = float(outlier_threshold) if not np.isfinite(outlier_threshold) or outlier_threshold <= 0: raise ValueError("outlier_threshold must be finite and positive") continuous = epochs_to_continuous(X) expected_weight_shape = (X.shape[0], X.shape[2]) if X.ndim == 3 else (X.shape[1],) if sample_weight is None: manual_weight = np.ones(continuous.shape[1], dtype=np.float64) else: sample_weight = np.asarray(sample_weight, dtype=np.float64) if sample_weight.shape != expected_weight_shape: raise ValueError( "sample_weight must have shape " f"{expected_weight_shape}, got {sample_weight.shape}" ) manual_weight = sample_weight.reshape(-1) if not np.isfinite(manual_weight).all() or np.any(manual_weight < 0): raise ValueError("sample_weight must be finite and non-negative") if np.count_nonzero(manual_weight > 0) < 2: raise ValueError( "sample_weight must weight at least two samples positively" ) automatic_weight = _automatic_sample_mask( continuous, manual_weight, outlier_threshold ) combined_weight = manual_weight * automatic_weight if np.count_nonzero(combined_weight > 0) < 2: raise ValueError( "fewer than two positively weighted samples remain after rejection" ) training_mean = (continuous @ combined_weight / combined_weight.sum())[ :, np.newaxis ] centered = continuous - training_mean current = centered composite = np.eye(continuous.shape[0], dtype=np.float64) matrices = [] ranks = [] effective_neighbors = 0 for iteration in range(int(n_iter)): cov = compute_covariance( current, weights=combined_weight, assume_centered=True, chunk_size=chunk_size, ) matrix, effective_neighbors, iteration_ranks = compute_sns_weights( cov, n_neighbors=n_neighbors, skip=skip, rcond=rcond ) matrices.append(matrix) ranks.append(iteration_ranks) composite = matrix @ composite if iteration + 1 < int(n_iter): current = apply_spatial_transform(matrix, current, chunk_size=chunk_size) cleaned = apply_spatial_transform(composite, centered, chunk_size=chunk_size) if preserve_mean: cleaned += training_mean cleaned = continuous_to_epochs(cleaned, X.shape) return cleaned, { "weights": composite, "denoising_matrix": composite, "denoising_matrices": tuple(matrices), "training_mean": training_mean, "n_neighbors": effective_neighbors, "requested_n_neighbors": int(n_neighbors), "skip": int(skip), "rcond": float(rcond), "preserve_mean": preserve_mean, "n_iter": int(n_iter), "outlier_threshold": outlier_threshold, "chunk_size": chunk_size, "neighbor_ranks": ranks[-1], "neighbor_ranks_per_iteration": tuple(ranks), "input_rank": int(np.linalg.matrix_rank(centered)), "effective_weight_sum": float(combined_weight.sum()), "rejected_sample_count": int(np.count_nonzero(automatic_weight == 0)), }
[docs] class SNS(BaseEstimator, TransformerMixin): """Sensor Noise Suppression estimator. The estimator fits a channel mean and one or more spatial projection operators on training data. Both are fixed during ``transform``. It accepts MNE Raw, Epochs, and Evoked objects or channel-first arrays and implements the SNS algorithm described in [1]_. Parameters ---------- n_neighbors : int, default=0 Neighbours used per channel. Zero uses all available neighbours. skip : int, default=0 Most-correlated neighbours to skip. rcond : float, default=1e-12 Relative pseudoinverse cutoff. preserve_mean : bool, default=False Restore the fitted training channel mean after regeneration. verbose : bool | str | int | None MNE-style logging level. n_iter : int, default=1 Number of successive SNS projections to learn and compose. outlier_threshold : float | None, default=None Maximum robust z-score retained while fitting. ``None`` disables automatic rejection. chunk_size : int | None, default=None Samples per chunk for statistics and operator application. MNE inputs are still materialized by the package's shared extractor. Attributes ---------- training_mean_ : ndarray, shape (n_channels, 1) Weighted channel mean learned during fit. denoising_matrix_ : ndarray, shape (n_channels, n_channels) Composite spatial operator. denoising_matrices_ : tuple of ndarray One spatial operator per iteration. neighbor_ranks_per_iteration_ : tuple of ndarray Local neighbor covariance ranks for every iteration. References ---------- .. [1] de Cheveigné, A., & Simon, J. Z. (2008). Sensor noise suppression. Journal of Neuroscience Methods, 168(1), 195-202. https://doi.org/10.1016/j.jneumeth.2007.09.012 """
[docs] def __init__( self, n_neighbors: int = 0, skip: int = 0, rcond: float = _DEFAULT_RCOND, preserve_mean: bool = False, verbose: bool | str | int | None = None, n_iter: int = 1, outlier_threshold: float | None = None, chunk_size: int | None = None, ) -> None: self.n_neighbors = n_neighbors self.skip = skip self.rcond = rcond self.preserve_mean = preserve_mean self.verbose = verbose self.n_iter = n_iter self.outlier_threshold = outlier_threshold self.chunk_size = chunk_size
def fit(self, X: Any, y=None, sample_weight: np.ndarray | None = None) -> SNS: """Learn fitted means and SNS operators from ``X``. Parameters ---------- X : array-like | mne.io.BaseRaw | mne.BaseEpochs | mne.Evoked Data used to learn the SNS operator. y : None Ignored. Included for scikit-learn compatibility. sample_weight : ndarray | None, default=None Non-negative fitting weights with shape ``(n_times,)`` for continuous data or ``(n_epochs, n_times)`` for epoched data. Returns ------- self : SNS Fitted estimator. """ set_log_level_from_verbose(self.verbose) data, _sfreq, _mne_type, _orig, _picks, names = extract_data_from_mne( X, auto_pick=True ) _cleaned, info = compute_sns( np.asarray(data, dtype=np.float64), n_neighbors=self.n_neighbors, skip=self.skip, rcond=self.rcond, preserve_mean=self.preserve_mean, n_iter=self.n_iter, outlier_threshold=self.outlier_threshold, chunk_size=self.chunk_size, sample_weight=sample_weight, ) self.training_mean_ = info["training_mean"] self.denoising_matrix_ = info["denoising_matrix"] self.denoising_matrices_ = info["denoising_matrices"] self.n_neighbors_ = info["n_neighbors"] self.neighbor_ranks_per_iteration_ = info["neighbor_ranks_per_iteration"] self.neighbor_ranks_ = self.neighbor_ranks_per_iteration_[-1] self.input_rank_ = info["input_rank"] self.n_channels_in_ = self.denoising_matrix_.shape[0] self.n_iter_ = info["n_iter"] self.chunk_size_ = info["chunk_size"] self.effective_weight_sum_ = info["effective_weight_sum"] self.rejected_sample_count_ = info["rejected_sample_count"] self.feature_names_in_ = None if names is None else tuple(names) logger.info( "SNS: learned %d iteration(s) on %d channels (%d neighbours each; " "%d samples rejected).", self.n_iter_, self.denoising_matrix_.shape[0], self.n_neighbors_, self.rejected_sample_count_, ) return self def transform(self, X: Any, y=None) -> Any: """Apply the fitted SNS operator. Parameters ---------- X : array-like | mne.io.BaseRaw | mne.BaseEpochs | mne.Evoked Data with the same channel layout used during fitting. y : None Ignored. Included for scikit-learn compatibility. Returns ------- X_clean : same type as X A copy with selected data channels replaced by their SNS result. """ check_is_fitted(self, ("denoising_matrix_", "training_mean_")) set_log_level_from_verbose(self.verbose) data, _sfreq, mne_type, orig_inst, picks, names = extract_data_from_mne( X, auto_pick=True ) data = check_channel_first_data(data, name="SNS") if not isinstance(self.preserve_mean, bool): raise TypeError("preserve_mean must be a bool") check_channel_layout( "SNS", n_channels=data.shape[-2], fitted_n_channels=self.n_channels_in_, ch_names=None if names is None else tuple(names), fitted_ch_names=self.feature_names_in_, ) continuous = epochs_to_continuous(data) cleaned = apply_spatial_transform( self.denoising_matrix_, continuous - self.training_mean_, chunk_size=self.chunk_size_, ) if self.preserve_mean: cleaned += self.training_mean_ cleaned = continuous_to_epochs(cleaned, data.shape) return reconstruct_mne_object(cleaned, orig_inst, mne_type, picks=picks) def fit_transform( self, X: Any, y=None, *, sample_weight: np.ndarray | None = None, **fit_params, ) -> Any: """Fit on ``X`` and apply the fitted operator. Parameters ---------- X : array-like | mne.io.BaseRaw | mne.BaseEpochs | mne.Evoked Data to fit and transform. y : None Ignored. Included for scikit-learn compatibility. sample_weight : ndarray | None, default=None Non-negative fitting weights. **fit_params : dict Reserved for scikit-learn compatibility. Returns ------- X_clean : same type as X Sensor-noise-suppressed data. """ if fit_params: unexpected = ", ".join(sorted(fit_params)) raise TypeError(f"Unexpected fit parameters: {unexpected}") return self.fit(X, y, sample_weight=sample_weight).transform(X)