Allow Xdawn to be used with rank-deficient data without reg

AGENTS.md

@@ -114,6 +114,14 @@ User-facing changes need a file `doc/changes/dev/<PR-number>.<type>.rst` (types:
 to `doc/changes/names.inc` (build fails otherwise) and are credited with `:newcontrib:` in their
 changelog entry instead of a plain name link.
 
+The `<PR-number>` for a not-yet-opened PR is one more than the highest number currently in use;
+issues and PRs share a single number sequence, so query the most recently created of either
+(the `issues` API endpoint includes PRs):
+```bash
+gh api "repos/mne-tools/mne-python/issues?state=all&per_page=1&sort=created&direction=desc" \
+    --jq '.[0].number'
+```
+
 ## Code conventions (beyond what ruff enforces)
 
 - Classes: `CamelCase`. Functions/variables: `snake_case`, no abbreviated names like `nsamples`.

doc/changes/dev/14182.newfeature.rst

@@ -0,0 +1 @@
+Added a ``rank`` parameter to :class:`mne.preprocessing.Xdawn` for using Xdawn with rank-deficient data without regularization, by `Eric Larson`_.

examples/preprocessing/xdawn_denoising.py

@@ -49,14 +49,15 @@
 
 raw.info["bads"] = ["MEG 2443"]  # set bad channels
 picks = pick_types(raw.info, meg=True, eeg=False, stim=False, eog=False, exclude="bads")
-# Epoching
+# Epoching, applying the SSP projectors that come with this dataset
+raw.apply_proj()
 epochs = Epochs(
     raw,
     events,
     event_id,
     tmin,
     tmax,
-    proj=False,
+    proj=True,
     picks=picks,
     baseline=None,
     preload=True,
@@ -69,12 +70,19 @@
 # %%
 # Now, we estimate a set of xDAWN filters for the epochs (which contain only
 # the ``vis_r`` class).
+#
+# Applying the three SSP projectors above makes the data rank deficient (302
+# instead of 305), so the generalized eigenvalue decomposition that xDAWN
+# relies on is ill-conditioned (and can fail outright) unless we tell it about
+# the rank of the data. Passing ``rank="info"`` restricts the decomposition to
+# the 302-dimensional principal subspace of the signal covariance and projects
+# the resulting filters and patterns back out to the 305 sensors.
 
 # Estimates signal covariance
 signal_cov = compute_raw_covariance(raw, picks=picks)
 
 # Xdawn instance
-xd = Xdawn(n_components=2, signal_cov=signal_cov)
+xd = Xdawn(n_components=2, signal_cov=signal_cov, rank="info")
 
 # Fit xdawn
 xd.fit(epochs)

mne/preprocessing/tests/test_xdawn.py

@@ -12,6 +12,7 @@
 from mne import (
     Epochs,
     EpochsArray,
+    compute_rank,
     compute_raw_covariance,
     create_info,
     pick_types,
@@ -230,10 +231,33 @@ def test_xdawn_regularization():
     xd.fit(epochs)
     xd = Xdawn(correct_overlap=False, reg="diagonal_fixed")
     xd.fit(epochs)
-    # XXX in principle this should maybe raise an error due to deficiency?
-    # xd = Xdawn(correct_overlap=False, reg=None)
-    # with pytest.raises(ValueError, match='Could not compute eigenvalues'):
-    #     xd.fit(epochs)
+    # Without regularization the GED is ill-conditioned, and depending on the LAPACK
+    # implementation it can fail outright. Using rank to restrict it to the principal
+    # subspace of the data makes it well posed on any implementation.
+    n_channels = len(epochs.ch_names)
+    rank = compute_rank(epochs, rank="info")["meg"]
+    assert rank == n_channels - len(epochs.info["projs"]) < n_channels
+    for use_rank in ("info", dict(meg=rank)):
+        xd = Xdawn(correct_overlap=False, reg=None, rank=use_rank)
+        xd.fit(epochs)
+        for eid in epochs.event_id:
+            # filters and patterns are restricted, but live in sensor space
+            assert xd.filters_[eid].shape == (rank, n_channels)
+            assert xd.patterns_[eid].shape == (rank, n_channels)
+            assert_allclose(
+                xd.filters_[eid] @ xd.patterns_[eid].T, np.eye(rank), atol=1e-8
+            )
+        # keeping all components round-trips the (rank-deficient) data
+        epochs_r = xd.apply(epochs, include=list(range(rank)))["cond2"]
+        assert_allclose(
+            epochs_r.get_data(copy=False),
+            epochs.get_data(copy=False),
+            atol=1e-8 * np.abs(epochs.get_data(copy=False)).max(),
+        )
+    # asking for more components than the rank allows
+    xd = Xdawn(n_components=rank + 1, correct_overlap=False, reg=None, rank="info")
+    with pytest.warns(RuntimeWarning, match="Rank restriction left"):
+        xd.fit(epochs)
 
 
 def test_XdawnTransformer():

mne/preprocessing/xdawn.py

@@ -6,12 +6,14 @@
 from scipy import linalg
 
 from .._fiff.pick import _pick_data_channels, pick_info
-from ..cov import Covariance, _regularized_covariance
+from ..cov import Covariance, _compute_rank_raw_array, _regularized_covariance
+from ..decoding._covs_ged import _handle_info_rank
+from ..decoding._ged import _get_restr_mat, _smart_ged
 from ..decoding.xdawn import XdawnTransformer
 from ..epochs import BaseEpochs
 from ..evoked import Evoked, EvokedArray
 from ..io import BaseRaw
-from ..utils import _check_option, logger, pinv
+from ..utils import _check_option, fill_doc, logger, pinv, warn
 
 
 def _construct_signal_from_epochs(epochs, events, sfreq, tmin):
@@ -106,6 +108,7 @@ def _fit_xdawn(
     sfreq=1.0,
     method_params=None,
     info=None,
+    rank="full",
 ):
     """Fit filters and coefs using Xdawn Algorithm.
 
@@ -140,31 +143,34 @@ def _fit_xdawn(
     sfreq : float
         Sampling frequency.  Only used if events is passed to correct for
         epochs overlap.
+    rank : None | 'full' | 'info' | dict
+        The rank of the data. If not ``'full'``, the covariances are restricted
+        to their ``rank``-dimensional principal subspace before the generalized
+        eigendecomposition, and the resulting filters and patterns are
+        projected back out to the full sensor space.
 
     Returns
     -------
-    filters : array, shape (n_channels, n_channels)
+    filters : array, shape (n_class * n_filters, n_channels)
         The Xdawn components used to decompose the data for each event type.
-        Each row corresponds to one component.
-    patterns : array, shape (n_channels, n_channels)
+        Each row corresponds to one component. ``n_filters`` is ``n_channels``
+        unless ``rank`` restricts the decomposition, in which case it is the
+        rank of the data.
+    patterns : array, shape (n_class * n_filters, n_channels)
         The Xdawn patterns used to restore the signals for each event type.
-    evokeds : array, shape (n_class, n_components, n_times)
+    evokeds : array, shape (n_class, n_channels, n_times)
         The independent evoked responses per condition.
     """
     if not isinstance(epochs_data, np.ndarray) or epochs_data.ndim != 3:
         raise ValueError("epochs_data must be 3D ndarray")
 
     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.
+    info, rank = _handle_info_rank(epochs_data, info, rank)
 
     # Retrieve or compute whitening covariance
     if signal_cov is None:
         signal_cov = _regularized_covariance(
-            np.hstack(epochs_data), reg, method_params, info, rank="full"
+            np.hstack(epochs_data), reg, method_params, info, rank=rank
         )
     elif isinstance(signal_cov, Covariance):
         signal_cov = signal_cov.data
@@ -176,6 +182,22 @@ def _fit_xdawn(
             "or an array of shape (n_chans, n_chans)"
         )
 
+    # Restriction to the rank-dimensional principal subspace of signal_cov,
+    # in which the generalized eigendecomposition below is well posed
+    if isinstance(rank, str) and rank == "full":
+        restr_mat = None
+    else:
+        if not isinstance(rank, dict):
+            rank = _compute_rank_raw_array(
+                np.hstack(epochs_data),
+                info,
+                rank=rank,
+                scalings=None,
+                log_ch_type="data",
+                on_few_samples="ignore",
+            )
+        restr_mat = _get_restr_mat(signal_cov, info, rank)
+
     # Get prototype events
     if events is not None:
         evokeds, toeplitzs = _least_square_evoked(epochs_data, events, tmin, sfreq)
@@ -191,11 +213,11 @@ def _fit_xdawn(
     for evo, toeplitz in zip(evokeds, toeplitzs):
         # Estimate covariance matrix of the prototype response
         evo = np.dot(evo, toeplitz)
-        evo_cov = _regularized_covariance(evo, reg, method_params, info, rank="full")
+        evo_cov = _regularized_covariance(evo, reg, method_params, info, rank=rank)
 
         # Fit spatial filters
         try:
-            evals, evecs = linalg.eigh(evo_cov, signal_cov)
+            evals, evecs = _smart_ged(evo_cov, signal_cov, restr_mat)
         except np.linalg.LinAlgError as exp:
             raise ValueError(
                 f"Could not compute eigenvalues, ensure proper regularization ({exp})"
@@ -212,6 +234,7 @@ def _fit_xdawn(
     return filters, patterns, evokeds
 
 
+@fill_doc
 class Xdawn(XdawnTransformer):
     """Implementation of the Xdawn Algorithm.
 
@@ -239,13 +262,22 @@ class Xdawn(XdawnTransformer):
         If float, shrinkage is used (0 <= shrinkage <= 1).
         For str options, ``reg`` will be passed as ``method`` to
         :func:`mne.compute_covariance`.
+    %(rank_full)s
+        If not ``'full'``, the covariances are restricted to their
+        ``rank``-dimensional principal subspace before computing the spatial
+        filters, which are then projected back out to the full sensor space.
+        This is useful for rank-deficient data, e.g., data with SSP projectors
+        applied or that has been Maxwell filtered.
+
+        .. versionadded:: 1.13
 
     Attributes
     ----------
     filters_ : dict of ndarray
         If fit, the Xdawn components used to decompose the data for each event
         type, else empty. For each event type, the filters are in the rows of
-        the corresponding array.
+        the corresponding array (``n_channels`` rows, or ``rank`` rows if
+        ``rank`` is not ``'full'``).
     patterns_ : dict of ndarray
         If fit, the Xdawn patterns used to restore the signals for each event
         type, else empty.
@@ -270,10 +302,18 @@ class Xdawn(XdawnTransformer):
     """
 
     def __init__(
-        self, n_components=2, signal_cov=None, correct_overlap="auto", reg=None
+        self,
+        n_components=2,
+        signal_cov=None,
+        correct_overlap="auto",
+        reg=None,
+        *,
+        rank="full",
     ):
         """Init."""
-        super().__init__(n_components=n_components, signal_cov=signal_cov, reg=reg)
+        super().__init__(
+            n_components=n_components, signal_cov=signal_cov, reg=reg, rank=rank
+        )
         self.correct_overlap = _check_option(
             "correct_overlap", correct_overlap, ["auto", True, False]
         )
@@ -322,7 +362,8 @@ def fit(self, epochs, y=None):
         self.correct_overlap_ = correct_overlap
 
         # Note: In this original version of Xdawn we compute and keep all
-        # components. The selection comes at transform().
+        # components (or all components of the rank-restricted subspace).
+        # The selection comes at transform().
         n_components = X.shape[1]
 
         # Main fitting function
@@ -337,11 +378,20 @@ def fit(self, epochs, y=None):
             sfreq=sfreq,
             method_params=self.cov_method_params,
             info=use_info,
+            rank=self.rank,
         )
 
+        # rank restriction can leave fewer filters than channels
+        n_filters = filters.shape[0] // len(evokeds)
+        if n_filters < min(self.n_components, X.shape[1]):
+            warn(
+                f"Rank restriction left {n_filters} components, which is fewer than "
+                f"n_components ({self.n_components}), consider using a larger rank."
+            )
+
         # Re-order filters and patterns according to event_id
-        filters = filters.reshape(-1, n_components, filters.shape[-1])
-        patterns = patterns.reshape(-1, n_components, patterns.shape[-1])
+        filters = filters.reshape(-1, n_filters, filters.shape[-1])
+        patterns = patterns.reshape(-1, n_filters, patterns.shape[-1])
         self.filters_, self.patterns_, self.evokeds_ = dict(), dict(), dict()
         idx = np.argsort([value for _, value in epochs.event_id.items()])
         for eid, this_filter, this_pattern, this_evo in zip(
@@ -425,7 +475,8 @@ def apply(self, inst, event_id=None, include=None, exclude=None):
         picks = _pick_data_channels(inst.info)
 
         # Define the components to keep
-        default_exclude = list(range(self.n_components, len(inst.ch_names)))
+        n_filters = len(next(iter(self.filters_.values())))
+        default_exclude = list(range(self.n_components, n_filters))
         if exclude is None:
             exclude = default_exclude
         else: