# Authors: The MNE-Python contributors. # License: BSD-3-Clause # Copyright the MNE-Python contributors. import numpy as np from scipy.linalg import eigh from sklearn.base import BaseEstimator, TransformerMixin from .._fiff.pick import _picks_to_idx from ..cov import Covariance, _regularized_covariance from ..defaults import _handle_default from ..filter import filter_data from ..rank import compute_rank from ..time_frequency import psd_array_welch from ..utils import ( _check_option, _time_mask, _validate_type, _verbose_safe_false, fill_doc, logger, ) @fill_doc class SSD(TransformerMixin, BaseEstimator): """ Signal decomposition using the Spatio-Spectral Decomposition (SSD). SSD seeks to maximize the power at a frequency band of interest while simultaneously minimizing it at the flanking (surrounding) frequency bins (considered noise). It extremizes the covariance matrices associated with signal and noise :footcite:`NikulinEtAl2011`. SSD can either be used as a dimensionality reduction method or a ‘denoised’ low rank factorization method :footcite:`HaufeEtAl2014b`. Parameters ---------- %(info_not_none)s Must match the input data. filt_params_signal : dict Filtering for the frequencies of interest. filt_params_noise : dict Filtering for the frequencies of non-interest. reg : float | str | None (default) Which covariance estimator to use. If not None (same as 'empirical'), allow regularization for covariance estimation. If float, shrinkage is used (0 <= shrinkage <= 1). For str options, reg will be passed to method :func:`mne.compute_covariance`. n_components : int | None (default None) The number of components to extract from the signal. If None, the number of components equal to the rank of the data are returned (see ``rank``). picks : array of int | None (default None) The indices of good channels. sort_by_spectral_ratio : bool (default True) If set to True, the components are sorted according to the spectral ratio. See Eq. (24) in :footcite:`NikulinEtAl2011`. return_filtered : bool (default False) If return_filtered is True, data is bandpassed and projected onto the SSD components. n_fft : int (default None) If sort_by_spectral_ratio is set to True, then the SSD sources will be sorted according to their spectral ratio which is calculated based on :func:`mne.time_frequency.psd_array_welch`. The n_fft parameter sets the length of FFT used. See :func:`mne.time_frequency.psd_array_welch` for more information. cov_method_params : dict | None (default None) As in :class:`mne.decoding.SPoC` The default is None. rank : None | dict | ‘info’ | ‘full’ As in :class:`mne.decoding.SPoC` This controls the rank computation that can be read from the measurement info or estimated from the data, which determines the maximum possible number of components. See Notes of :func:`mne.compute_rank` for details. We recommend to use 'full' when working with epoched data. Attributes ---------- filters_ : array, shape (n_channels, n_components) The spatial filters to be multiplied with the signal. patterns_ : array, shape (n_components, n_channels) The patterns for reconstructing the signal from the filtered data. References ---------- .. footbibliography:: """ def __init__( self, info, filt_params_signal, filt_params_noise, reg=None, n_components=None, picks=None, sort_by_spectral_ratio=True, return_filtered=False, n_fft=None, cov_method_params=None, rank=None, ): """Initialize instance.""" dicts = {"signal": filt_params_signal, "noise": filt_params_noise} for param, dd in [("l", 0), ("h", 0), ("l", 1), ("h", 1)]: key = ("signal", "noise")[dd] if param + "_freq" not in dicts[key]: raise ValueError( f"{param + '_freq'} must be defined in filter parameters for {key}" ) val = dicts[key][param + "_freq"] if not isinstance(val, int | float): _validate_type(val, ("numeric",), f"{key} {param}_freq") # check freq bands if ( filt_params_noise["l_freq"] > filt_params_signal["l_freq"] or filt_params_signal["h_freq"] > filt_params_noise["h_freq"] ): raise ValueError( "Wrongly specified frequency bands!\n" "The signal band-pass must be within the noise " "band-pass!" ) self.picks = picks del picks self.info = info self.freqs_signal = (filt_params_signal["l_freq"], filt_params_signal["h_freq"]) self.freqs_noise = (filt_params_noise["l_freq"], filt_params_noise["h_freq"]) self.filt_params_signal = filt_params_signal self.filt_params_noise = filt_params_noise # check if boolean if not isinstance(sort_by_spectral_ratio, (bool)): raise ValueError("sort_by_spectral_ratio must be boolean") self.sort_by_spectral_ratio = sort_by_spectral_ratio if n_fft is None: self.n_fft = int(self.info["sfreq"]) else: self.n_fft = int(n_fft) # check if boolean if not isinstance(return_filtered, (bool)): raise ValueError("return_filtered must be boolean") self.return_filtered = return_filtered self.reg = reg self.n_components = n_components self.rank = rank self.cov_method_params = cov_method_params def _check_X(self, X): """Check input data.""" _validate_type(X, np.ndarray, "X") _check_option("X.ndim", X.ndim, (2, 3)) n_chan = X.shape[-2] if n_chan != self.info["nchan"]: raise ValueError( "Info must match the input data." f"Found {n_chan} channels but expected {self.info['nchan']}." ) def fit(self, X, y=None): """Estimate the SSD decomposition on raw or epoched data. Parameters ---------- X : array, shape ([n_epochs, ]n_channels, n_times) The input data from which to estimate the SSD. Either 2D array obtained from continuous data or 3D array obtained from epoched data. y : None Ignored; exists for compatibility with scikit-learn pipelines. Returns ------- self : instance of SSD Returns the modified instance. """ ch_types = self.info.get_channel_types(picks=self.picks, unique=True) if len(ch_types) > 1: raise ValueError( "At this point SSD only supports fitting " f"single channel types. Your info has {len(ch_types)} types." ) self.picks_ = _picks_to_idx(self.info, self.picks, none="data", exclude="bads") self._check_X(X) X_aux = X[..., self.picks_, :] X_signal = filter_data(X_aux, self.info["sfreq"], **self.filt_params_signal) X_noise = filter_data(X_aux, self.info["sfreq"], **self.filt_params_noise) X_noise -= X_signal if X.ndim == 3: X_signal = np.hstack(X_signal) X_noise = np.hstack(X_noise) # prevent rank change when computing cov with rank='full' cov_signal = _regularized_covariance( X_signal, reg=self.reg, method_params=self.cov_method_params, rank="full", info=self.info, ) cov_noise = _regularized_covariance( X_noise, reg=self.reg, method_params=self.cov_method_params, rank="full", info=self.info, ) # project cov to rank subspace cov_signal, cov_noise, rank_proj = _dimensionality_reduction( cov_signal, cov_noise, self.info, self.rank ) eigvals_, eigvects_ = eigh(cov_signal, cov_noise) # sort in descending order ix = np.argsort(eigvals_)[::-1] self.eigvals_ = eigvals_[ix] # project back to sensor space self.filters_ = np.matmul(rank_proj, eigvects_[:, ix]) self.patterns_ = np.linalg.pinv(self.filters_) # We assume that ordering by spectral ratio is more important # than the initial ordering. This ordering should be also learned when # fitting. X_ssd = self.filters_.T @ X[..., self.picks_, :] sorter_spec = Ellipsis if self.sort_by_spectral_ratio: _, sorter_spec = self.get_spectral_ratio(ssd_sources=X_ssd) self.sorter_spec = sorter_spec logger.info("Done.") return self def transform(self, X): """Estimate epochs sources given the SSD filters. Parameters ---------- X : array, shape ([n_epochs, ]n_channels, n_times) The input data from which to estimate the SSD. Either 2D array obtained from continuous data or 3D array obtained from epoched data. Returns ------- X_ssd : array, shape ([n_epochs, ]n_components, n_times) The processed data. """ self._check_X(X) if self.filters_ is None: raise RuntimeError("No filters available. Please first call fit") if self.return_filtered: X_aux = X[..., self.picks_, :] X = filter_data(X_aux, self.info["sfreq"], **self.filt_params_signal) X_ssd = self.filters_.T @ X[..., self.picks_, :] if X.ndim == 2: X_ssd = X_ssd[self.sorter_spec][: self.n_components] else: X_ssd = X_ssd[:, self.sorter_spec, :][:, : self.n_components, :] return X_ssd def fit_transform(self, X, y=None, **fit_params): """Fit SSD to data, then transform it. Fits transformer to ``X`` and ``y`` with optional parameters ``fit_params``, and returns a transformed version of ``X``. Parameters ---------- X : array, shape ([n_epochs, ]n_channels, n_times) The input data from which to estimate the SSD. Either 2D array obtained from continuous data or 3D array obtained from epoched data. y : None Ignored; exists for compatibility with scikit-learn pipelines. **fit_params : dict Additional fitting parameters passed to the :meth:`mne.decoding.SSD.fit` method. Not used for this class. Returns ------- X_ssd : array, shape ([n_epochs, ]n_components, n_times) The processed data. """ # use parent TransformerMixin method but with custom docstring return super().fit_transform(X, y=y, **fit_params) def get_spectral_ratio(self, ssd_sources): """Get the spectal signal-to-noise ratio for each spatial filter. Spectral ratio measure for best n_components selection See :footcite:`NikulinEtAl2011`, Eq. (24). Parameters ---------- ssd_sources : array Data projected to SSD space. Returns ------- spec_ratio : array, shape (n_channels) Array with the sprectal ratio value for each component. sorter_spec : array, shape (n_channels) Array of indices for sorting spec_ratio. References ---------- .. footbibliography:: """ psd, freqs = psd_array_welch( ssd_sources, sfreq=self.info["sfreq"], n_fft=self.n_fft ) sig_idx = _time_mask(freqs, *self.freqs_signal) noise_idx = _time_mask(freqs, *self.freqs_noise) if psd.ndim == 3: mean_sig = psd[:, :, sig_idx].mean(axis=2).mean(axis=0) mean_noise = psd[:, :, noise_idx].mean(axis=2).mean(axis=0) spec_ratio = mean_sig / mean_noise else: mean_sig = psd[:, sig_idx].mean(axis=1) mean_noise = psd[:, noise_idx].mean(axis=1) spec_ratio = mean_sig / mean_noise sorter_spec = spec_ratio.argsort()[::-1] return spec_ratio, sorter_spec def inverse_transform(self): """Not implemented yet.""" raise NotImplementedError("inverse_transform is not yet available.") def apply(self, X): """Remove selected components from the signal. This procedure will reconstruct M/EEG signals from which the dynamics described by the excluded components is subtracted (denoised by low-rank factorization). See :footcite:`HaufeEtAl2014b` for more information. .. note:: Unlike in other classes with an apply method, only NumPy arrays are supported (not instances of MNE objects). Parameters ---------- X : array, shape ([n_epochs, ]n_channels, n_times) The input data from which to estimate the SSD. Either 2D array obtained from continuous data or 3D array obtained from epoched data. Returns ------- X : array, shape ([n_epochs, ]n_channels, n_times) The processed data. """ X_ssd = self.transform(X) pick_patterns = self.patterns_[self.sorter_spec][: self.n_components].T X = pick_patterns @ X_ssd return X def _dimensionality_reduction(cov_signal, cov_noise, info, rank): """Perform dimensionality reduction on the covariance matrices.""" n_channels = cov_signal.shape[0] # find ranks of covariance matrices rank_signal = list( compute_rank( Covariance( cov_signal, info.ch_names, list(), list(), 0, verbose=_verbose_safe_false(), ), rank, _handle_default("scalings_cov_rank", None), info, ).values() )[0] rank_noise = list( compute_rank( Covariance( cov_noise, info.ch_names, list(), list(), 0, verbose=_verbose_safe_false(), ), rank, _handle_default("scalings_cov_rank", None), info, ).values() )[0] rank = np.min([rank_signal, rank_noise]) # should be identical if rank < n_channels: eigvals, eigvects = eigh(cov_signal) # sort in descending order ix = np.argsort(eigvals)[::-1] eigvals = eigvals[ix] eigvects = eigvects[:, ix] # compute rank subspace projection matrix rank_proj = np.matmul( eigvects[:, :rank], np.eye(rank) * (eigvals[:rank] ** -0.5) ) logger.info( "Projecting covariance of %i channels to %i rank subspace", n_channels, rank, ) else: rank_proj = np.eye(n_channels) logger.info("Preserving covariance rank (%i)", rank) # project covariance matrices to rank subspace cov_signal = np.matmul(rank_proj.T, np.matmul(cov_signal, rank_proj)) cov_noise = np.matmul(rank_proj.T, np.matmul(cov_noise, rank_proj)) return cov_signal, cov_noise, rank_proj