BUG: Fix bug with CSP rank (#12476)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: larsoner

doc/changes/devel/12476.bugfix.rst

@@ -0,0 +1 @@
+Fixed bugs with handling of rank in :class:`mne.decoding.CSP`, by `Eric Larson`_.

examples/decoding/decoding_csp_eeg.py

@@ -49,6 +49,7 @@
 montage = make_standard_montage("standard_1005")
 raw.set_montage(montage)
 raw.annotations.rename(dict(T1="hands", T2="feet"))
+raw.set_eeg_reference(projection=True)
 
 # Apply band-pass filter
 raw.filter(7.0, 30.0, fir_design="firwin", skip_by_annotation="edge")

mne/cov.py

@@ -59,7 +59,7 @@
     empirical_covariance,
     log_likelihood,
 )
-from .rank import compute_rank
+from .rank import _compute_rank
 from .utils import (
     _array_repr,
     _check_fname,
@@ -1226,6 +1226,21 @@ def _eigvec_subspace(eig, eigvec, mask):
     return eig, eigvec
 
 
+@verbose
+def _compute_rank_raw_array(
+    data, info, rank, scalings, *, log_ch_type=None, verbose=None
+):
+    from .io import RawArray
+
+    return _compute_rank(
+        RawArray(data, info, copy=None, verbose=_verbose_safe_false()),
+        rank,
+        scalings,
+        info,
+        log_ch_type=log_ch_type,
+    )
+
+
 def _compute_covariance_auto(
     data,
     method,
@@ -1237,22 +1252,31 @@ def _compute_covariance_auto(
     stop_early,
     picks_list,
     rank,
+    *,
+    cov_kind="",
+    log_ch_type=None,
+    log_rank=True,
 ):
     """Compute covariance auto mode."""
-    from .io import RawArray
-
     # rescale to improve numerical stability
     orig_rank = rank
-    rank = compute_rank(
-        RawArray(data.T, info, copy=None, verbose=_verbose_safe_false()),
-        rank,
-        scalings,
+    rank = _compute_rank_raw_array(
+        data.T,
         info,
+        rank=rank,
+        scalings=scalings,
+        verbose=_verbose_safe_false(),
     )
     with _scaled_array(data.T, picks_list, scalings):
         C = np.dot(data.T, data)
         _, eigvec, mask = _smart_eigh(
-            C, info, rank, proj_subspace=True, do_compute_rank=False
+            C,
+            info,
+            rank,
+            proj_subspace=True,
+            do_compute_rank=False,
+            log_ch_type=log_ch_type,
+            verbose=None if log_rank else _verbose_safe_false(),
         )
         eigvec = eigvec[mask]
         data = np.dot(data, eigvec.T)
@@ -1261,21 +1285,24 @@ def _compute_covariance_auto(
             (key, np.searchsorted(used, picks)) for key, picks in picks_list
         ]
         sub_info = pick_info(info, used) if len(used) != len(mask) else info
-        logger.info(f"Reducing data rank from {len(mask)} -> {eigvec.shape[0]}")
+        if log_rank:
+            logger.info(f"Reducing data rank from {len(mask)} -> {eigvec.shape[0]}")
         estimator_cov_info = list()
-        msg = "Estimating covariance using {}"
 
         ok_sklearn = check_version("sklearn")
         if not ok_sklearn and (len(method) != 1 or method[0] != "empirical"):
             raise ValueError(
-                "scikit-learn is not installed, `method` must be `empirical`, got "
-                f"{method}"
+                'scikit-learn is not installed, `method` must be "empirical", got '
+                f"{repr(method)}"
             )
 
         for method_ in method:
             data_ = data.copy()
             name = method_.__name__ if callable(method_) else method_
-            logger.info(msg.format(name.upper()))
+            logger.info(
+                f'Estimating {cov_kind + (" " if cov_kind else "")}'
+                f"covariance using {name.upper()}"
+            )
             mp = method_params[method_]
             _info = {}
 
@@ -1691,9 +1718,8 @@ def _get_ch_whitener(A, pca, ch_type, rank):
     mask[:-rank] = False
 
     logger.info(
-        "    Setting small {} eigenvalues to zero ({})".format(
-            ch_type, "using PCA" if pca else "without PCA"
-        )
+        f"    Setting small {ch_type} eigenvalues to zero "
+        f'({"using" if pca else "without"} PCA)'
     )
     if pca:  # No PCA case.
         # This line will reduce the actual number of variables in data
@@ -1791,6 +1817,8 @@ def _smart_eigh(
     proj_subspace=False,
     do_compute_rank=True,
     on_rank_mismatch="ignore",
+    *,
+    log_ch_type=None,
     verbose=None,
 ):
     """Compute eigh of C taking into account rank and ch_type scalings."""
@@ -1813,8 +1841,13 @@ def _smart_eigh(
 
     noise_cov = Covariance(C, ch_names, [], projs, 0)
     if do_compute_rank:  # if necessary
-        rank = compute_rank(
-            noise_cov, rank, scalings, info, on_rank_mismatch=on_rank_mismatch
+        rank = _compute_rank(
+            noise_cov,
+            rank,
+            scalings,
+            info,
+            on_rank_mismatch=on_rank_mismatch,
+            log_ch_type=log_ch_type,
         )
     assert C.ndim == 2 and C.shape[0] == C.shape[1]
 
@@ -1838,7 +1871,11 @@ def _smart_eigh(
         else:
             this_rank = rank[ch_type]
 
-        e, ev, m = _get_ch_whitener(this_C, False, ch_type.upper(), this_rank)
+        if log_ch_type is not None:
+            ch_type_ = log_ch_type
+        else:
+            ch_type_ = ch_type.upper()
+        e, ev, m = _get_ch_whitener(this_C, False, ch_type_, this_rank)
         if proj_subspace:
             # Choose the subspace the same way we do for projections
             e, ev = _eigvec_subspace(e, ev, m)
@@ -1995,7 +2032,7 @@ def regularize(
     else:
         regs.update(mag=mag, grad=grad)
     if rank != "full":
-        rank = compute_rank(cov, rank, scalings, info)
+        rank = _compute_rank(cov, rank, scalings, info)
 
     info_ch_names = info["ch_names"]
     ch_names_by_type = dict()
@@ -2071,7 +2108,17 @@ def regularize(
     return cov
 
 
-def _regularized_covariance(data, reg=None, method_params=None, info=None, rank=None):
+def _regularized_covariance(
+    data,
+    reg=None,
+    method_params=None,
+    info=None,
+    rank=None,
+    *,
+    log_ch_type=None,
+    log_rank=None,
+    cov_kind="",
+):
     """Compute a regularized covariance from data using sklearn.
 
     This is a convenience wrapper for mne.decoding functions, which
@@ -2114,6 +2161,9 @@ def _regularized_covariance(data, reg=None, method_params=None, info=None, rank=
         picks_list=picks_list,
         scalings=scalings,
         rank=rank,
+        cov_kind=cov_kind,
+        log_ch_type=log_ch_type,
+        log_rank=log_rank,
     )[reg]["data"]
     return cov
 

mne/decoding/csp.py

@@ -12,11 +12,18 @@
 import numpy as np
 from scipy.linalg import eigh
 
-from ..cov import _regularized_covariance
+from .._fiff.meas_info import 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
 from ..fixes import pinv
-from ..utils import _check_option, _validate_type, copy_doc, fill_doc
+from ..utils import (
+    _check_option,
+    _validate_type,
+    _verbose_safe_false,
+    copy_doc,
+    fill_doc,
+)
 from .base import BaseEstimator
 from .mixin import TransformerMixin
 
@@ -185,6 +192,9 @@ def fit(self, X, y):
                 f"{n_classes} classes; use component_order='mutual_info' instead."
             )
 
+        # Convert rank to one that will run
+        _validate_type(self.rank, (dict, None), "rank")
+
         covs, sample_weights = self._compute_covariance_matrices(X, y)
         eigen_vectors, eigen_values = self._decompose_covs(covs, sample_weights)
         ix = self._order_components(
@@ -519,10 +529,28 @@ def _compute_covariance_matrices(self, X, y):
         elif self.cov_est == "epoch":
             cov_estimator = self._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
+        self._info = create_info(n_channels, 1000.0, "mag")
+        if self.rank is None:
+            self._rank = _compute_rank_raw_array(
+                X.transpose(1, 0, 2).reshape(X.shape[1], -1),
+                self._info,
+                rank=None,
+                scalings=None,
+                log_ch_type="data",
+            )
+        else:
+            self._rank = {"mag": sum(self.rank.values())}
+
         covs = []
         sample_weights = []
-        for this_class in self._classes:
-            cov, weight = cov_estimator(X[y == this_class])
+        for ci, this_class in enumerate(self._classes):
+            cov, weight = cov_estimator(
+                X[y == this_class],
+                cov_kind=f"class={this_class}",
+                log_rank=ci == 0,
+            )
 
             if self.norm_trace:
                 cov /= np.trace(cov)
@@ -532,29 +560,39 @@ def _compute_covariance_matrices(self, X, y):
 
         return np.stack(covs), np.array(sample_weights)
 
-    def _concat_cov(self, x_class):
+    def _concat_cov(self, x_class, *, cov_kind, log_rank):
         """Concatenate epochs before computing the covariance."""
         _, n_channels, _ = x_class.shape
 
-        x_class = np.transpose(x_class, [1, 0, 2])
-        x_class = x_class.reshape(n_channels, -1)
+        x_class = x_class.transpose(1, 0, 2).reshape(n_channels, -1)
         cov = _regularized_covariance(
-            x_class, reg=self.reg, method_params=self.cov_method_params, rank=self.rank
+            x_class,
+            reg=self.reg,
+            method_params=self.cov_method_params,
+            rank=self._rank,
+            info=self._info,
+            cov_kind=cov_kind,
+            log_rank=log_rank,
+            log_ch_type="data",
         )
         weight = x_class.shape[0]
 
         return cov, weight
 
-    def _epoch_cov(self, x_class):
+    def _epoch_cov(self, x_class, *, cov_kind, log_rank):
         """Mean of per-epoch covariances."""
         cov = sum(
             _regularized_covariance(
                 this_X,
                 reg=self.reg,
                 method_params=self.cov_method_params,
-                rank=self.rank,
+                rank=self._rank,
+                info=self._info,
+                cov_kind=cov_kind,
+                log_rank=log_rank and ii == 0,
+                log_ch_type="data",
             )
-            for this_X in x_class
+            for ii, this_X in enumerate(x_class)
         )
         cov /= len(x_class)
         weight = len(x_class)
@@ -563,6 +601,20 @@ def _epoch_cov(self, x_class):
 
     def _decompose_covs(self, covs, sample_weights):
         n_classes = len(covs)
+        n_channels = covs[0].shape[0]
+        assert self._rank is not None  # should happen in _compute_covariance_matrices
+        _, sub_vec, mask = _smart_eigh(
+            covs.mean(0),
+            self._info,
+            self._rank,
+            proj_subspace=True,
+            do_compute_rank=False,
+            log_ch_type="data",
+            verbose=_verbose_safe_false(),
+        )
+        sub_vec = sub_vec[mask]
+        covs = np.array([sub_vec @ cov @ sub_vec.T for cov in covs], float)
+        assert covs[0].shape == (mask.sum(),) * 2
         if n_classes == 2:
             eigen_values, eigen_vectors = eigh(covs[0], covs.sum(0))
         else:
@@ -573,6 +625,9 @@ def _decompose_covs(self, covs, sample_weights):
                 eigen_vectors.T, covs, sample_weights
             )
             eigen_values = None
+        # project back
+        eigen_vectors = sub_vec.T @ eigen_vectors
+        assert eigen_vectors.shape == (n_channels, mask.sum())
         return eigen_vectors, eigen_values
 
     def _compute_mutual_info(self, covs, sample_weights, eigen_vectors):
@@ -824,6 +879,8 @@ def fit(self, X, y):
                 reg=self.reg,
                 method_params=self.cov_method_params,
                 rank=self.rank,
+                log_ch_type="data",
+                log_rank=ii == 0,
             )
 
         C = covs.mean(0)