Add big sklearn compliance test

Author: Genuster

mne/decoding/base.py

@@ -8,22 +8,23 @@
 import numbers
 
 import numpy as np
+import scipy.linalg
 from sklearn import model_selection as models
 from sklearn.base import (  # noqa: F401
     BaseEstimator,
     MetaEstimatorMixin,
-    TransformerMixin,
     clone,
     is_classifier,
 )
 from sklearn.linear_model import LogisticRegression
 from sklearn.metrics import check_scoring
 from sklearn.model_selection import KFold, StratifiedKFold, check_cv
 from sklearn.utils import check_array, check_X_y, indexable
+from sklearn.utils.validation import check_is_fitted
 
 from ..parallel import parallel_func
 from ..utils import _pl, logger, pinv, verbose, warn
-from .ged import _get_ssd_rank, _handle_restr_map, _smart_ajd, _smart_ged
+from .ged import _handle_restr_map, _smart_ajd, _smart_ged
 from .transformer import MNETransformerMixin
 
 
@@ -55,7 +56,7 @@ class GEDTransformer(MNETransformerMixin, BaseEstimator):
         (except the last) returned by cov_callable is decomposed with the last
         covariance. In this case, number of covariances should be number of classes + 1.
         Defaults to "single".
-    restr_map : "restricting" | "whitening" | "ssd" | None
+    restr_type : "restricting" | "whitening" | "ssd" | None
         Restricting transformation for covariance matrices before performing GED.
         If "restricting" only restriction to the principal subspace of the C_ref
         will be performed.
@@ -94,7 +95,7 @@ def __init__(
         mod_ged_callable,
         mod_params,
         dec_type="single",
-        restr_map=None,
+        restr_type=None,
         R_func=None,
     ):
         self.n_filters = n_filters
@@ -103,27 +104,35 @@ def __init__(
         self.mod_ged_callable = mod_ged_callable
         self.mod_params = mod_params
         self.dec_type = dec_type
-        self.restr_map = restr_map
+        self.restr_type = restr_type
         self.R_func = R_func
 
     def fit(self, X, y=None):
         """..."""
+        X, y = self._check_data(
+            X,
+            y=y,
+            fit=True,
+            return_y=True,
+            atleast_3d=False if self.restr_type == "ssd" else True,
+        )
         covs, C_ref, info, rank, kwargs = self.cov_callable(X, y, **self.cov_params)
+        self._validate_covariances(covs + [C_ref])
         if self.dec_type == "single":
             if len(covs) > 2:
+                covs = np.array(covs)
                 sample_weights = kwargs["sample_weights"]
-                restr_map = _handle_restr_map(C_ref, self.restr_map, info, rank)
+                restr_map = _handle_restr_map(C_ref, self.restr_type, info, rank)
                 evecs = _smart_ajd(covs, restr_map, weights=sample_weights)
                 evals = None
             else:
                 S = covs[0]
                 R = covs[1]
-                if self.restr_map == "ssd":
-                    rank = _get_ssd_rank(S, R, info, rank)
+                if self.restr_type == "ssd":
                     mult_order = "ssd"
                 else:
                     mult_order = None
-                restr_map = _handle_restr_map(C_ref, self.restr_map, info, rank)
+                restr_map = _handle_restr_map(C_ref, self.restr_type, info, rank)
                 evals, evecs = _smart_ged(
                     S, R, restr_map, R_func=self.R_func, mult_order=mult_order
                 )
@@ -133,19 +142,26 @@ def fit(self, X, y=None):
             )
             self.evals_ = evals
             self.filters_ = evecs.T
-            if self.restr_map == "ssd":
+            if self.restr_type == "ssd":
                 self.patterns_ = np.linalg.pinv(evecs)
             else:
                 self.patterns_ = pinv(evecs)
 
         elif self.dec_type == "multi":
             self.classes_ = np.unique(y)
             R = covs[-1]
-            restr_map = _handle_restr_map(C_ref, self.restr_map, info, rank)
+            if self.restr_type == "ssd":
+                mult_order = "ssd"
+            else:
+                mult_order = None
+            restr_map = _handle_restr_map(C_ref, self.restr_type, info, rank)
             all_evals, all_evecs, all_patterns = list(), list(), list()
             for i in range(len(self.classes_)):
                 S = covs[i]
-                evals, evecs = _smart_ged(S, R, restr_map, R_func=self.R_func)
+
+                evals, evecs = _smart_ged(
+                    S, R, restr_map, R_func=self.R_func, mult_order=mult_order
+                )
 
                 evals, evecs = self.mod_ged_callable(
                     evals, evecs, covs, **self.mod_params, **kwargs
@@ -161,9 +177,48 @@ def fit(self, X, y=None):
 
     def transform(self, X):
         """..."""
-        X = np.dot(self.filters_, X)
+        check_is_fitted(self, "filters_")
+        X = self._check_data(X)
+        if self.dec_type == "single":
+            pick_filters = self.filters_[: self.n_filters]
+        elif self.dec_type == "multi":
+            pick_filters = np.concatenate(
+                [
+                    self.filters_[i, : self.n_filters]
+                    for i in range(self.filters_.shape[0])
+                ],
+                axis=0,
+            )
+        X = np.asarray([pick_filters @ epoch for epoch in X])
         return X
 
+    def _validate_covariances(self, covs):
+        for cov in covs:
+            if cov is None:
+                continue
+            is_sym = scipy.linalg.issymmetric(cov, rtol=1e-10, atol=1e-11)
+            if not is_sym:
+                raise ValueError(
+                    "One of covariances or C_ref is not symmetric, "
+                    "check your cov_callable"
+                )
+            if not np.all(np.linalg.eigvals(cov) >= 0):
+                ValueError(
+                    "One of covariances or C_ref has negative eigenvalues, "
+                    "check your cov_callable"
+                )
+
+    def __sklearn_tags__(self):
+        """Tag the transformer."""
+        tags = super().__sklearn_tags__()
+        tags.estimator_type = "transformer"
+        # Can be a transformer where S and R covs are not based on y classes.
+        tags.target_tags.required = False
+        tags.target_tags.one_d_labels = True
+        tags.input_tags.two_d_array = True
+        tags.input_tags.three_d_array = True
+        return tags
+
 
 class LinearModel(MetaEstimatorMixin, BaseEstimator):
     """Compute and store patterns from linear models.

mne/decoding/covs_ged.py

@@ -10,7 +10,9 @@
 from .._fiff.meas_info import Info, create_info
 from .._fiff.pick import _picks_to_idx
 from ..cov import Covariance, _compute_rank_raw_array, _regularized_covariance
+from ..defaults import _handle_default
 from ..filter import filter_data
+from ..rank import compute_rank
 from ..utils import _verbose_safe_false, logger, pinv
 
 
@@ -293,6 +295,26 @@ def _ssd_estimate(
     )
     covs = [S, R]
     C_ref = S
+
+    all_ranks = list()
+    for cov in covs:
+        r = list(
+            compute_rank(
+                Covariance(
+                    cov,
+                    info.ch_names,
+                    list(),
+                    list(),
+                    0,
+                    verbose=_verbose_safe_false(),
+                ),
+                rank,
+                _handle_default("scalings_cov_rank", None),
+                info,
+            ).values()
+        )[0]
+    all_ranks.append(r)
+    rank = np.min(all_ranks)
     return covs, C_ref, info, rank, dict()
 
 

mne/decoding/csp.py

@@ -142,7 +142,7 @@ def __init__(
             _csp_mod,
             mod_params,
             dec_type="single",
-            restr_map="restricting",
+            restr_type="restricting",
             R_func=sum,
         )
 
@@ -911,7 +911,7 @@ def __init__(
             _spoc_mod,
             mod_params,
             dec_type="single",
-            restr_map=None,
+            restr_type=None,
             R_func=None,
         )
 

mne/decoding/ged.py

@@ -6,28 +6,26 @@
 import scipy.linalg
 
 from ..cov import Covariance, _smart_eigh, compute_whitener
-from ..defaults import _handle_default
-from ..rank import compute_rank
-from ..utils import _verbose_safe_false, logger
+from ..utils import logger
 
 
-def _handle_restr_map(C_ref, restr_map, info, rank):
+def _handle_restr_map(C_ref, restr_type, info, rank):
     """Get restricting map to C_ref rank-dimensional principal subspace.
 
     Returns matrix of shape (rank, n_chs) used to restrict or
     restrict+rescale (whiten) covariances matrices.
     """
-    if C_ref is None or restr_map is None:
+    if C_ref is None or restr_type is None:
         return None
-    if restr_map == "whitening":
+    if restr_type == "whitening":
         projs = info["projs"]
         C_ref_cov = Covariance(C_ref, info.ch_names, info["bads"], projs, 0)
-        restr_map = compute_whitener(C_ref_cov, info, rank=rank, pca=True)
-    elif restr_map == "ssd":
+        restr_map = compute_whitener(C_ref_cov, info, rank=rank, pca=True)[0]
+    elif restr_type == "ssd":
         restr_map = _get_ssd_whitener(C_ref, rank)
-    elif restr_map == "restricting":
+    elif restr_type == "restricting":
         restr_map = _get_restricting_map(C_ref, info, rank)
-    elif isinstance(restr_map, callable):
+    elif isinstance(restr_type, callable):
         pass
     else:
         raise ValueError(
@@ -147,6 +145,15 @@ def _ajd_pham(X, eps=1e-6, max_iter=15):
     return V, D
 
 
+def _is_all_pos_def(covs):
+    for cov in covs:
+        try:
+            _ = scipy.linalg.cholesky(cov)
+        except np.linalg.LinAlgError:
+            return False
+    return True
+
+
 def _smart_ajd(covs, restr_map=None, weights=None):
     """Perform smart approximate joint diagonalization.
 
@@ -157,6 +164,12 @@ def _smart_ajd(covs, restr_map=None, weights=None):
     The matrix of generalized eigenvectors is of shape (n_chs, r).
     """
     if restr_map is None:
+        is_all_pos_def = _is_all_pos_def(covs)
+        if not is_all_pos_def:
+            raise ValueError(
+                "If C_ref is not provided by covariance estimator, "
+                "all the covs should be positive definite"
+            )
         evecs, D = _ajd_pham(covs)
         return evecs
 
@@ -191,42 +204,6 @@ def _normalize_eigenvectors(evecs, covs, sample_weights):
     return evecs
 
 
-def _get_ssd_rank(S, R, info, rank):
-    # find ranks of covariance matrices
-    rank_signal = list(
-        compute_rank(
-            Covariance(
-                S,
-                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(
-                R,
-                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
-    return rank
-
-
 def _get_ssd_whitener(S, rank):
     """Perform dimensionality reduction on the covariance matrices."""
     n_channels = S.shape[0]

mne/decoding/ssd.py

@@ -137,7 +137,7 @@ def __init__(
             _ssd_mod,
             mod_params,
             dec_type="single",
-            restr_map="ssd",
+            restr_type="ssd",
             R_func=None,
         )