add option for CSP to select restr_type and provide info

Author: Genuster

mne/decoding/_covs_ged.py

@@ -15,7 +15,7 @@
 from ..utils import _verbose_safe_false, logger
 
 
-def _concat_cov(x_class, *, cov_kind, log_rank, reg, cov_method_params, rank, info):
+def _concat_cov(x_class, *, cov_kind, log_rank, reg, cov_method_params, info, rank):
     """Concatenate epochs before computing the covariance."""
     _, n_channels, _ = x_class.shape
 
@@ -34,7 +34,7 @@ def _concat_cov(x_class, *, cov_kind, log_rank, reg, cov_method_params, rank, in
     return cov, n_channels  # the weight here is just the number of channels
 
 
-def _epoch_cov(x_class, *, cov_kind, log_rank, reg, cov_method_params, rank, info):
+def _epoch_cov(x_class, *, cov_kind, log_rank, reg, cov_method_params, info, rank):
     """Mean of per-epoch covariances."""
     name = reg if isinstance(reg, str) else "empirical"
     name += " with shrinkage" if isinstance(reg, float) else ""
@@ -62,22 +62,29 @@ def _epoch_cov(x_class, *, cov_kind, log_rank, reg, cov_method_params, rank, inf
     return cov, weight
 
 
-def _csp_estimate(X, y, reg, cov_method_params, cov_est, rank, norm_trace):
+def _handle_info_rank(X, info, rank):
+    if info is None:
+        # use mag instead of eeg to avoid the cov EEG projection warning
+        info = create_info(X.shape[1], 1000.0, "mag")
+        if isinstance(rank, dict):
+            rank = dict(mag=sum(rank.values()))
+
+    return info, rank
+
+
+def _csp_estimate(X, y, reg, cov_method_params, cov_est, info, rank, norm_trace):
     _, n_channels, _ = X.shape
     classes_ = np.unique(y)
     if cov_est == "concat":
         cov_estimator = _concat_cov
     elif cov_est == "epoch":
         cov_estimator = _epoch_cov
-    # Someday we could allow the user to pass this, then we wouldn't need to convert
-    # but in the meantime they can use a pipeline with a scaler
-    _info = create_info(n_channels, 1000.0, "mag")
-    if isinstance(rank, dict):
-        _rank = {"mag": sum(rank.values())}
-    else:
-        _rank = _compute_rank_raw_array(
-            X.transpose(1, 0, 2).reshape(X.shape[1], -1),
-            _info,
+
+    info, rank = _handle_info_rank(X, info, rank)
+    if not isinstance(rank, dict):
+        rank = _compute_rank_raw_array(
+            np.hstack(X),
+            info,
             rank=rank,
             scalings=None,
             log_ch_type="data",
@@ -92,8 +99,8 @@ def _csp_estimate(X, y, reg, cov_method_params, cov_est, rank, norm_trace):
             log_rank=ci == 0,
             reg=reg,
             cov_method_params=cov_method_params,
-            rank=_rank,
-            info=_info,
+            info=info,
+            rank=rank,
         )
 
         if norm_trace:
@@ -105,7 +112,7 @@ def _csp_estimate(X, y, reg, cov_method_params, cov_est, rank, norm_trace):
     covs = np.stack(covs)
     C_ref = covs.mean(0)
 
-    return covs, C_ref, _info, _rank, dict(sample_weights=np.array(sample_weights))
+    return covs, C_ref, info, rank, dict(sample_weights=np.array(sample_weights))
 
 
 def _xdawn_estimate(
@@ -118,6 +125,7 @@ def _xdawn_estimate(
     rank="full",
 ):
     classes = np.unique(y)
+    info, rank = _handle_info_rank(X, info, rank)
 
     # Retrieve or compute whitening covariance
     if R is None:
@@ -143,7 +151,14 @@ def _xdawn_estimate(
 
     covs.append(R)
     C_ref = R
-    rank = rank if isinstance(rank, dict) else None
+    if not isinstance(rank, dict):
+        rank = _compute_rank_raw_array(
+            np.hstack(X),
+            info,
+            rank=rank,
+            scalings=None,
+            log_ch_type="data",
+        )
     return covs, C_ref, info, rank, dict()
 
 

mne/decoding/base.py

@@ -22,8 +22,6 @@
 from sklearn.utils import check_array, check_X_y, indexable
 from sklearn.utils.validation import check_is_fitted
 
-from .._fiff.meas_info import create_info
-from ..cov import _compute_rank_raw_array
 from ..parallel import parallel_func
 from ..utils import _pl, logger, pinv, verbose, warn
 from ._ged import _handle_restr_mat, _is_cov_symm_pos_semidef, _smart_ajd, _smart_ged
@@ -70,7 +68,7 @@ class _GEDTransformer(MNETransformerMixin, BaseEstimator):
         preserved for compatibility.
         If None, no restriction will be applied. Defaults to None.
     R_func : callable | None
-        If provided GED will be performed on (S, R_func(S,R)).
+        If provided, GED will be performed on (S, R_func(S,R)).
 
     Attributes
     ----------
@@ -88,7 +86,10 @@ class _GEDTransformer(MNETransformerMixin, BaseEstimator):
     CSP
     SPoC
     SSD
-    mne.preprocessing.Xdawn
+
+    Notes
+    -----
+    .. versionadded:: 1.10
     """
 
     def __init__(
@@ -125,23 +126,6 @@ def fit(self, X, y=None):
             self.mod_ged_callable if self.mod_ged_callable is not None else _no_op_mod
         )
 
-        # If restriction to be done, info and rank should exist.
-        if self.restr_type is not None and C_ref is not None:
-            if info is None:
-                # use mag instead of eeg to avoid the cov EEG projection warning
-                info = create_info(C_ref.shape[0], 1000.0, "mag")
-                if isinstance(rank, dict):
-                    rank = dict(mag=sum(rank.values()))
-
-            if rank is None:
-                rank = _compute_rank_raw_array(
-                    np.hstack(X),
-                    info,
-                    rank=None,
-                    scalings=None,
-                    log_ch_type="data",
-                )
-
         if self.dec_type == "single":
             if len(covs) > 2:
                 weights = (

mne/decoding/csp.py

@@ -9,7 +9,7 @@
 from scipy.linalg import eigh
 from sklearn.utils.validation import check_is_fitted
 
-from .._fiff.meas_info import create_info
+from .._fiff.meas_info import Info, create_info
 from ..cov import _compute_rank_raw_array, _regularized_covariance, _smart_eigh
 from ..defaults import _BORDER_DEFAULT, _EXTRAPOLATE_DEFAULT, _INTERPOLATION_DEFAULT
 from ..evoked import EvokedArray
@@ -70,6 +70,26 @@ class CSP(_GEDTransformer):
         Parameters to pass to :func:`mne.compute_covariance`.
 
         .. versionadded:: 0.16
+
+    restr_type : "restricting" | "whitening" | None
+        Restricting transformation for covariance matrices before performing
+        generalized eigendecomposition.
+        If "restricting" only restriction to the principal subspace of signal_cov
+        will be performed.
+        If "whitening", covariance matrices will be additionally rescaled according
+        to the whitening for the signal_cov.
+        If None, no restriction will be applied. Defaults to "restricting".
+
+        .. versionadded:: 1.10
+    info : mne.Info | None
+        The mne.Info object with information about the sensors and methods of
+        measurement used for covariance estimation and generalized
+        eigendecomposition.
+        If None, one channel type and no projections will be assumed and if
+        rank is dict, it will be sum of ranks per channel type.
+        Defaults to None.
+
+        .. versionadded:: 1.10
     %(rank_none)s
 
         .. versionadded:: 0.17
@@ -113,11 +133,14 @@ def __init__(
         transform_into="average_power",
         norm_trace=False,
         cov_method_params=None,
+        restr_type="restricting",
+        info=None,
         rank=None,
         component_order="mutual_info",
     ):
         # Init default CSP
         self.n_components = n_components
+        self.info = info
         self.rank = rank
         self.reg = reg
         self.cov_est = cov_est
@@ -126,12 +149,14 @@ def __init__(
         self.norm_trace = norm_trace
         self.cov_method_params = cov_method_params
         self.component_order = component_order
+        self.restr_type = restr_type
 
         cov_callable = partial(
             _csp_estimate,
             reg=reg,
             cov_method_params=cov_method_params,
             cov_est=cov_est,
+            info=info,
             rank=rank,
             norm_trace=norm_trace,
         )
@@ -140,7 +165,7 @@ def __init__(
             n_components=n_components,
             cov_callable=cov_callable,
             mod_ged_callable=mod_ged_callable,
-            restr_type="restricting",
+            restr_type=restr_type,
             R_func=sum,
         )
 
@@ -172,6 +197,12 @@ def _validate_params(self, *, y):
         n_classes = len(self.classes_)
         if n_classes < 2:
             raise ValueError(f"n_classes must be >= 2, but got {n_classes} class")
+        _check_option(
+            "restr_type",
+            self.restr_type,
+            ("restricting", "whitening", None),
+        )
+        _validate_type(self.info, (Info, None), "info")
 
     def fit(self, X, y):
         """Estimate the CSP decomposition on epochs.