ENH: Add native save and read support for SSD (#13718)

Co-authored-by: Thomas S. Binns <t.s.binns@outlook.com>

Author: Aniketsy

doc/api/decoding.rst

@@ -41,3 +41,7 @@ Functions that assist with decoding and model fitting:
    cross_val_multiscore
    get_coef
    get_spatial_filter_from_estimator
+   read_csp
+   read_spoc
+   read_ssd
+   read_xdawn_transformer

doc/changes/dev/13718.other.rst

@@ -0,0 +1 @@
+Add native save and read functionality for :class:`mne.decoding.SSD`, :class:`mne.decoding.CSP`, :class:`mne.decoding.SPoC`, and :class:`mne.decoding.XdawnTransformer` objects, by `Aniket Singh Yadav`_.

mne/decoding/__init__.pyi

@@ -10,6 +10,10 @@ __all__ = [
     "SPoC",
     "SSD",
     "Scaler",
+    "read_csp",
+    "read_spoc",
+    "read_ssd",
+    "read_xdawn_transformer",
     "SlidingEstimator",
     "SpatialFilter",
     "TemporalFilter",
@@ -31,12 +35,12 @@ from .base import (
     cross_val_multiscore,
     get_coef,
 )
-from .csp import CSP, SPoC
+from .csp import CSP, SPoC, read_csp, read_spoc
 from .ems import EMS, compute_ems
 from .receptive_field import ReceptiveField
 from .search_light import GeneralizingEstimator, SlidingEstimator
 from .spatial_filter import SpatialFilter, get_spatial_filter_from_estimator
-from .ssd import SSD
+from .ssd import SSD, read_ssd
 from .time_delaying_ridge import TimeDelayingRidge
 from .time_frequency import TimeFrequency
 from .transformer import (
@@ -47,4 +51,4 @@ from .transformer import (
     UnsupervisedSpatialFilter,
     Vectorizer,
 )
-from .xdawn import XdawnTransformer
+from .xdawn import XdawnTransformer, read_xdawn_transformer

mne/decoding/base.py

@@ -25,16 +25,20 @@
 from sklearn.utils import indexable
 from sklearn.utils.validation import check_is_fitted
 
+from .._fiff.meas_info import Info
 from ..parallel import parallel_func
 from ..utils import (
+    _check_fname,
     _check_option,
+    _import_h5io_funcs,
     _pl,
     _validate_type,
     logger,
     pinv,
     verbose,
     warn,
 )
+from ..utils.check import check_fname
 from ._fixes import validate_data
 from ._ged import (
     _handle_restr_mat,
@@ -135,6 +139,65 @@ def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         cls._is_base_ged = False
 
+    def __getstate__(self):
+        """Prepare state for serialization."""
+        state = self.__dict__.copy()
+        state.pop("cov_callable", None)
+        state.pop("mod_ged_callable", None)
+        state.pop("R_func", None)
+        return state
+
+    def _restore_callables(self):
+        """Restore callables after loading serialized state."""
+        # Expected to be implemented in child classes
+        pass
+
+    def __setstate__(self, state):
+        """Restore state from serialization."""
+        required_state_keys = getattr(self, "_required_state_keys", ())
+        missing = [k for k in required_state_keys if k not in state]
+        if missing:
+            raise ValueError(
+                f"Cannot read file as type {type(self).__name__}, "
+                f"it is missing required keys: {missing}. "
+                "Please report this to MNE developers "
+                "(https://github.com/mne-tools/mne-python/issues/new) "
+                "and include a copy of this entire traceback and a link "
+                "to the problematic file with your report."
+            )
+        if "info" in state and isinstance(state["info"], dict):
+            state["info"] = Info(**state["info"])
+        self.__dict__.update(state)
+        self._restore_callables()
+
+    @verbose
+    def save(self, fname, *, overwrite=False, verbose=None):
+        """Save the object to disk (in HDF5 format).
+
+        Parameters
+        ----------
+        fname : path-like
+            The file path to save to. Should end with ``'.h5'`` or
+            ``'.hdf5'``.
+        %(overwrite)s
+        %(verbose)s
+
+        Notes
+        -----
+        .. versionadded:: 1.12
+        """
+        _, write_hdf5 = _import_h5io_funcs()
+        class_name = getattr(self, "_save_fname_type", type(self).__name__.lower())
+        check_fname(fname, class_name, (".h5", ".hdf5"))
+        fname = _check_fname(fname, overwrite=overwrite, verbose=verbose)
+        write_hdf5(
+            fname,
+            self.__getstate__(),
+            overwrite=overwrite,
+            title="mnepython",
+            slash="replace",
+        )
+
     def fit(self, X, y=None):
         """..."""
         # Let the inheriting transformers check data by themselves
@@ -338,6 +401,18 @@ def __sklearn_tags__(self):
         return tags
 
 
+@verbose
+def _read_ged(fname, ged_class, *, verbose=None):
+    """Load a saved GED transformer object from disk."""
+    read_hdf5, _ = _import_h5io_funcs()
+    _validate_type(fname, "path-like", "fname")
+    fname = _check_fname(fname, overwrite=True, must_exist=True, verbose=verbose)
+    state = read_hdf5(fname, title="mnepython", slash="replace")
+    inst = object.__new__(ged_class)
+    inst.__setstate__(state)
+    return inst
+
+
 class LinearModel(MetaEstimatorMixin, BaseEstimator):
     """Compute and store patterns from linear models.
 

mne/decoding/csp.py

@@ -9,10 +9,16 @@
 
 from .._fiff.meas_info import Info
 from ..defaults import _BORDER_DEFAULT, _EXTRAPOLATE_DEFAULT, _INTERPOLATION_DEFAULT
-from ..utils import _check_option, _validate_type, fill_doc, legacy
+from ..utils import (
+    _check_option,
+    _validate_type,
+    fill_doc,
+    legacy,
+    verbose,
+)
 from ._covs_ged import _csp_estimate, _spoc_estimate
 from ._mod_ged import _csp_mod, _spoc_mod
-from .base import _GEDTransformer
+from .base import _GEDTransformer, _read_ged
 from .spatial_filter import get_spatial_filter_from_estimator
 
 
@@ -160,13 +166,43 @@ def __init__(
             R_func=sum,
         )
 
+    _save_fname_type = "csp"
+
     def __sklearn_tags__(self):
         """Tag the transformer."""
         tags = super().__sklearn_tags__()
         tags.target_tags.required = True
         tags.target_tags.multi_output = True
         return tags
 
+    _required_state_keys = (
+        "component_order",
+        "cov_est",
+        "cov_method_params",
+        "info",
+        "log",
+        "n_components",
+        "norm_trace",
+        "rank",
+        "reg",
+        "restr_type",
+        "transform_into",
+    )
+
+    def _restore_callables(self):
+        """Restore CSP-specific callables after loading state."""
+        self.cov_callable = partial(
+            _csp_estimate,
+            reg=self.reg,
+            cov_method_params=self.cov_method_params,
+            cov_est=self.cov_est,
+            info=self.info,
+            rank=self.rank,
+            norm_trace=self.norm_trace,
+        )
+        self.mod_ged_callable = partial(_csp_mod, evecs_order=self.component_order)
+        self.R_func = sum
+
     def _validate_params(self, *, y):
         _validate_type(self.n_components, int, "n_components")
         if hasattr(self, "cov_est"):
@@ -766,12 +802,37 @@ def __init__(
         delattr(self, "cov_est")
         delattr(self, "norm_trace")
 
+    _save_fname_type = "spoc"
+
     def __sklearn_tags__(self):
         """Tag the transformer."""
         tags = super().__sklearn_tags__()
         tags.target_tags.multi_output = False
         return tags
 
+    _required_state_keys = (
+        "cov_method_params",
+        "info",
+        "log",
+        "n_components",
+        "rank",
+        "reg",
+        "restr_type",
+        "transform_into",
+    )
+
+    def _restore_callables(self):
+        """Restore SPoC-specific callables after loading state."""
+        self.cov_callable = partial(
+            _spoc_estimate,
+            reg=self.reg,
+            cov_method_params=self.cov_method_params,
+            info=self.info,
+            rank=self.rank,
+        )
+        self.mod_ged_callable = _spoc_mod
+        self.R_func = None
+
     def fit(self, X, y):
         """Estimate the SPoC decomposition on epochs.
 
@@ -848,3 +909,57 @@ def fit_transform(self, X, y=None, **fit_params):
         """
         # use parent TransformerMixin method but with custom docstring
         return super().fit_transform(X, y=y, **fit_params)
+
+
+@verbose
+def read_csp(fname, *, verbose=None):
+    """Load a saved :class:`mne.decoding.CSP` object from disk.
+
+    Parameters
+    ----------
+    fname : path-like
+        Path to a CSP file in HDF5 format, which should end with ``.h5`` or
+        ``.hdf5``.
+    %(verbose)s
+
+    Returns
+    -------
+    csp : instance of :class:`~mne.decoding.CSP`
+        The loaded CSP object with all fitted attributes restored.
+
+    See Also
+    --------
+    mne.decoding.CSP.save
+
+    Notes
+    -----
+    .. versionadded:: 1.12
+    """
+    return _read_ged(fname, CSP, verbose=verbose)
+
+
+@verbose
+def read_spoc(fname, *, verbose=None):
+    """Load a saved :class:`mne.decoding.SPoC` object from disk.
+
+    Parameters
+    ----------
+    fname : path-like
+        Path to a SPoC file in HDF5 format, which should end with ``.h5`` or
+        ``.hdf5``.
+    %(verbose)s
+
+    Returns
+    -------
+    spoc : instance of :class:`~mne.decoding.SPoC`
+        The loaded SPoC object with all fitted attributes restored.
+
+    See Also
+    --------
+    mne.decoding.SPoC.save
+
+    Notes
+    -----
+    .. versionadded:: 1.12
+    """
+    return _read_ged(fname, SPoC, verbose=verbose)