address Eric's suggestions

Author: Genuster

mne/decoding/ged.py mne/decoding/_ged.pymne/decoding/ged.py renamed to mne/decoding/_ged.py

@@ -24,7 +24,7 @@
 
 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
+from ._ged import _handle_restr_mat, _is_cov_symm_pos_semidef, _smart_ajd, _smart_ged
 from .transformer import MNETransformerMixin
 
 
@@ -40,9 +40,9 @@ class _GEDTransformer(MNETransformerMixin, BaseEstimator):
         The number of spatial filters to decompose M/EEG signals.
     cov_callable : callable
         Function used to estimate covariances and reference matrix (C_ref) from the
-        data.
-    cov_params : dict
-        Parameters passed to cov_callable.
+        data. It should accept only X and y as arguments and return covs, C_ref, info,
+        rank and additional kwargs passed further to mod_ged_callable.
+        C_ref, info, rank can be None, while kwargs can be empty dict.
     mod_ged_callable : callable
         Function used to modify (e.g. sort or normalize) generalized
         eigenvalues and eigenvectors.
@@ -91,7 +91,6 @@ def __init__(
         self,
         n_components,
         cov_callable,
-        cov_params,
         mod_ged_callable,
         *,
         mod_params=None,
@@ -101,7 +100,6 @@ def __init__(
     ):
         self.n_components = n_components
         self.cov_callable = cov_callable
-        self.cov_params = cov_params
         self.mod_ged_callable = mod_ged_callable
         self.mod_params = mod_params
         self.dec_type = dec_type
@@ -117,7 +115,7 @@ def fit(self, X, y=None):
             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)
+        covs, C_ref, info, rank, kwargs = self.cov_callable(X, y)
         covs = np.stack(covs)
         self._validate_covariances(covs)
         self._validate_covariances([C_ref])

mne/decoding/base.py

@@ -3,6 +3,7 @@
 # Copyright the MNE-Python contributors.
 
 import copy as cp
+from functools import partial
 
 import numpy as np
 from scipy.linalg import eigh
@@ -126,21 +127,19 @@ def __init__(
         self.cov_method_params = cov_method_params
         self.component_order = component_order
 
-        cov_params = dict(
+        cov_callable = partial(
+            _csp_estimate,
             reg=reg,
             cov_method_params=cov_method_params,
             cov_est=cov_est,
             rank=rank,
             norm_trace=norm_trace,
         )
-
         super().__init__(
             n_components=n_components,
-            cov_callable=_csp_estimate,
-            cov_params=cov_params,
+            cov_callable=cov_callable,
             mod_ged_callable=_csp_mod,
             mod_params=dict(evecs_order=component_order),
-            dec_type="single",
             restr_type="restricting",
             R_func=sum,
         )
@@ -899,20 +898,16 @@ def __init__(
             cov_method_params=cov_method_params,
         )
 
-        cov_params = dict(
+        cov_callable = partial(
+            _spoc_estimate,
             reg=reg,
             cov_method_params=cov_method_params,
             rank=rank,
         )
-
         super(CSP, self).__init__(
-            n_components,
-            _spoc_estimate,
-            cov_params,
-            _spoc_mod,
-            dec_type="single",
-            restr_type=None,
-            R_func=None,
+            n_components=n_components,
+            cov_callable=cov_callable,
+            mod_ged_callable=_spoc_mod,
         )
 
         # Covariance estimation have to be done on the single epoch level,

mne/decoding/csp.py

@@ -2,6 +2,8 @@
 # License: BSD-3-Clause
 # Copyright the MNE-Python contributors.
 
+from functools import partial
+
 import numpy as np
 from scipy.linalg import eigh
 from sklearn.utils.validation import check_is_fitted
@@ -119,7 +121,8 @@ def __init__(
         self.cov_method_params = cov_method_params
         self.rank = rank
 
-        cov_params = dict(
+        cov_callable = partial(
+            _ssd_estimate,
             reg=reg,
             cov_method_params=cov_method_params,
             info=info,
@@ -128,17 +131,11 @@ def __init__(
             filt_params_noise=filt_params_noise,
             rank=rank,
         )
-
-        mod_params = dict()
         super().__init__(
-            n_components,
-            _ssd_estimate,
-            cov_params,
-            _ssd_mod,
-            mod_params,
-            dec_type="single",
+            n_components=n_components,
+            cov_callable=cov_callable,
+            mod_ged_callable=_ssd_mod,
             restr_type="ssd",
-            R_func=None,
         )
 
     def _validate_params(self, X):

mne/decoding/ssd.py

@@ -2,7 +2,7 @@
 # License: BSD-3-Clause
 # Copyright the MNE-Python contributors.
 
-import functools
+from functools import partial
 from pathlib import Path
 
 import numpy as np
@@ -18,15 +18,15 @@
 from mne import Epochs, compute_rank, create_info, pick_types, read_events
 from mne._fiff.proj import make_eeg_average_ref_proj
 from mne.cov import Covariance, _regularized_covariance
-from mne.decoding.base import _GEDTransformer
-from mne.decoding.ged import (
+from mne.decoding._ged import (
     _get_restr_mat,
     _handle_restr_mat,
     _is_cov_pos_def,
     _is_cov_symm_pos_semidef,
     _smart_ajd,
     _smart_ged,
 )
+from mne.decoding.base import _GEDTransformer
 from mne.io import read_raw
 
 data_dir = Path(__file__).parents[2] / "io" / "tests" / "data"
@@ -120,20 +120,16 @@ def _mock_mod_ged_callable(evals, evecs, covs, **kwargs):
 
 param_grid = dict(
     n_components=[4],
-    cov_callable=[_mock_cov_callable],
-    cov_params=[
-        dict(cov_method_params=dict(reg="empirical")),
-    ],
+    cov_callable=[partial(_mock_cov_callable, cov_method_params=dict(reg="empirical"))],
     mod_ged_callable=[_mock_mod_ged_callable],
-    mod_params=[dict()],
     dec_type=["single", "multi"],
     # XXX: Not covering "ssd" here because test_ssd.py works with 2D data.
     # Need to fix its tests first.
     restr_type=[
         "restricting",
         "whitening",
     ],
-    R_func=[functools.partial(np.sum, axis=0)],
+    R_func=[partial(np.sum, axis=0)],
 )
 
 ged_estimators = [_GEDTransformer(**p) for p in ParameterGrid(param_grid)]
@@ -185,11 +181,8 @@ def test_ged_binary_cov():
     ged = _GEDTransformer(
         n_components=4,
         cov_callable=_mock_cov_callable,
-        cov_params=dict(),
         mod_ged_callable=_mock_mod_ged_callable,
-        dec_type="single",
         restr_type="restricting",
-        R_func=None,
     )
     ged.fit(X, y)
     desired_evals = ged.evals_
@@ -212,11 +205,9 @@ def test_ged_binary_cov():
     ged = _GEDTransformer(
         n_components=4,
         cov_callable=_mock_cov_callable,
-        cov_params=dict(),
         mod_ged_callable=_mock_mod_ged_callable,
         dec_type="multi",
         restr_type="restricting",
-        R_func=None,
     )
     ged.fit(X, y)
     desired_evals = ged.evals_
@@ -241,11 +232,8 @@ def test_ged_multicov():
     ged = _GEDTransformer(
         n_components=4,
         cov_callable=_mock_cov_callable,
-        cov_params=dict(),
         mod_ged_callable=_mock_mod_ged_callable,
-        dec_type="single",
         restr_type="restricting",
-        R_func=None,
     )
     ged.fit(X, y)
     desired_filters = ged.filters_
@@ -267,11 +255,9 @@ def test_ged_multicov():
     ged = _GEDTransformer(
         n_components=4,
         cov_callable=_mock_cov_callable,
-        cov_params=dict(),
         mod_ged_callable=_mock_mod_ged_callable,
         dec_type="multi",
         restr_type="restricting",
-        R_func=None,
     )
     ged.fit(X, y)
     desired_evals = ged.evals_
@@ -292,12 +278,11 @@ def test_ged_multicov():
 
     ged = _GEDTransformer(
         n_components=4,
-        cov_callable=_mock_cov_callable,
-        cov_params=dict(cov_method_params=dict(reg="oas"), compute_C_ref=False),
+        cov_callable=partial(
+            _mock_cov_callable, cov_method_params=dict(reg="oas"), compute_C_ref=False
+        ),
         mod_ged_callable=_mock_mod_ged_callable,
-        dec_type="single",
         restr_type="restricting",
-        R_func=None,
     )
     ged.fit(X, y)
     desired_filters = ged.filters_
@@ -310,11 +295,7 @@ def test_ged_invalid_cov():
     ged = _GEDTransformer(
         n_components=1,
         cov_callable=_mock_cov_callable,
-        cov_params=dict(),
         mod_ged_callable=_mock_mod_ged_callable,
-        dec_type="single",
-        restr_type=None,
-        R_func=None,
     )
     asymm_cov = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
     with pytest.raises(ValueError, match="not symmetric"):

mne/decoding/tests/test_ged.py

@@ -3,6 +3,7 @@
 # Copyright the MNE-Python contributors.
 
 from collections.abc import Mapping
+from functools import partial
 
 import numpy as np
 from scipy import linalg
@@ -263,16 +264,14 @@ def __init__(self, n_components=2, reg=None, signal_cov=None, method_params=None
         self.reg = reg
         self.method_params = method_params
 
-        cov_params = dict(reg=reg, cov_method_params=method_params, R=signal_cov)
-
+        cov_callable = partial(
+            _xdawn_estimate, reg=reg, cov_method_params=method_params, R=signal_cov
+        )
         super().__init__(
-            n_components,
-            _xdawn_estimate,
-            cov_params,
-            _xdawn_mod,
+            n_components=n_components,
+            cov_callable=cov_callable,
+            mod_ged_callable=_xdawn_mod,
             dec_type="multi",
-            restr_type=None,
-            R_func=None,
         )
 
     def _validate_params(self, X):