add feature to perform GED in the principal subspace for xdawn

Author: Genuster

mne/decoding/_covs_ged.py

@@ -119,11 +119,6 @@ def _xdawn_estimate(
 ):
     classes = np.unique(y)
 
-    # XXX Eventually this could be made to deal with rank deficiency properly
-    # by exposing this "rank" parameter, but this will require refactoring
-    # the linalg.eigh call to operate in the lower-dimension
-    # subspace, then project back out.
-
     # Retrieve or compute whitening covariance
     if R is None:
         R = _regularized_covariance(
@@ -147,9 +142,8 @@ def _xdawn_estimate(
         covs.append(evo_cov)
 
     covs.append(R)
-    C_ref = None
-    rank = None
-    info = None
+    C_ref = R
+    rank = rank if isinstance(rank, dict) else None
     return covs, C_ref, info, rank, dict()
 
 

mne/decoding/_ged.py

@@ -18,8 +18,7 @@ def _handle_restr_mat(C_ref, restr_type, info, rank):
     if C_ref is None or restr_type is None:
         return None
     if restr_type == "whitening":
-        projs = info["projs"]
-        C_ref_cov = Covariance(C_ref, info.ch_names, info["bads"], projs, 0)
+        C_ref_cov = Covariance(C_ref, info.ch_names, info["bads"], info["projs"], 0)
         restr_mat = compute_whitener(C_ref_cov, info, rank=rank, pca=True)[0]
     elif restr_type == "ssd":
         restr_mat = _get_ssd_whitener(C_ref, rank)

mne/decoding/base.py

@@ -22,6 +22,8 @@
 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
@@ -122,6 +124,24 @@ def fit(self, X, y=None):
         mod_ged_callable = (
             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/preprocessing/xdawn.py

@@ -8,6 +8,7 @@
 import numpy as np
 from scipy import linalg
 
+from .._fiff.meas_info import Info
 from .._fiff.pick import _pick_data_channels, pick_info
 from ..cov import Covariance, _regularized_covariance
 from ..decoding._covs_ged import _xdawn_estimate
@@ -246,6 +247,30 @@ class _XdawnTransformer(_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 None.
+
+        .. 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)s
+        Defaults to "full".
+
+        .. versionadded:: 1.10
+
 
     Attributes
     ----------
@@ -257,21 +282,39 @@ class _XdawnTransformer(_GEDTransformer):
         The Xdawn patterns used to restore the signals for each event type.
     """
 
-    def __init__(self, n_components=2, reg=None, signal_cov=None, method_params=None):
+    def __init__(
+        self,
+        n_components=2,
+        reg=None,
+        signal_cov=None,
+        method_params=None,
+        restr_type=None,
+        info=None,
+        rank="full",
+    ):
         """Init."""
         self.n_components = n_components
         self.signal_cov = signal_cov
         self.reg = reg
         self.method_params = method_params
+        self.restr_type = restr_type
+        self.info = info
+        self.rank = rank
 
         cov_callable = partial(
-            _xdawn_estimate, reg=reg, cov_method_params=method_params, R=signal_cov
+            _xdawn_estimate,
+            reg=reg,
+            cov_method_params=method_params,
+            R=signal_cov,
+            info=info,
+            rank=rank,
         )
         super().__init__(
             n_components=n_components,
             cov_callable=cov_callable,
             mod_ged_callable=_xdawn_mod,
             dec_type="multi",
+            restr_type=restr_type,
         )
 
     def _validate_params(self, X):
@@ -288,7 +331,13 @@ def _validate_params(self, X):
                 raise ValueError(
                     "signal_cov data should be of shape (n_channels, n_channels)"
                 )
-        _validate_type(self.method_params, (Mapping, None))
+        _validate_type(self.method_params, (Mapping, None), "method_params")
+        _check_option(
+            "restr_type",
+            self.restr_type,
+            ("restricting", "whitening", None),
+        )
+        _validate_type(self.info, (Info, None), "info")
 
     def fit(self, X, y=None):
         """Fit Xdawn spatial filters.