ENH: Cache bad channel interpolation matrix for EEG (Part 1) (#4651)

doc/api/preprocessing.rst

@@ -34,6 +34,8 @@ Projections:
 
    Layout
    DigMontage
+   InterpolationMatrix
+   compute_bad_channel_interpolation_matrix
    compute_native_head_t
    fix_mag_coil_types
    read_polhemus_fastscan

doc/changes/dev/14455.newfeature.rst

@@ -0,0 +1 @@
+Add :func:`mne.channels.compute_bad_channel_interpolation_matrix` and support for precomputed interpolation matrices in :meth:`~mne.io.Raw.interpolate_bads`, :meth:`~mne.Epochs.interpolate_bads`, and :meth:`~mne.Evoked.interpolate_bads`, by `Ishaan Maheshwari`_.

doc/conf.py

@@ -255,6 +255,7 @@
     "SourceSpaces": "mne.SourceSpaces",
     "Epochs": "mne.Epochs",
     "Layout": "mne.channels.Layout",
+    "InterpolationMatrix": "mne.channels.InterpolationMatrix",
     "EvokedArray": "mne.EvokedArray",
     "BiHemiLabel": "mne.BiHemiLabel",
     "AverageTFR": "mne.time_frequency.AverageTFR",

mne/channels/__init__.pyi

@@ -1,5 +1,7 @@
 __all__ = [
     "DigMontage",
+    "InterpolationMatrix",
+    "compute_bad_channel_interpolation_matrix",
     "Layout",
     "_EEG_SELECTIONS",
     "_SELECTIONS",
@@ -52,6 +54,10 @@ from .channels import (
     rename_channels,
     unify_bad_channels,
 )
+from .interpolation import (
+    InterpolationMatrix,
+    compute_bad_channel_interpolation_matrix,
+)
 from .layout import (
     Layout,
     find_layout,

mne/channels/channels.py

@@ -1058,6 +1058,8 @@ def interpolate_bads(
         method=None,
         exclude=(),
         on_bad_position="warn",
+        *,
+        interpolation=None,
         verbose=None,
     ):
         """Interpolate bad MEG and EEG channels.
@@ -1109,6 +1111,14 @@ def interpolate_bads(
             filled with :data:`~numpy.nan`.
 
             .. versionadded:: 1.12
+        interpolation : dict | instance of InterpolationMatrix | None
+            A precomputed interpolation matrix, e.g. from
+            :func:`mne.channels.compute_bad_channel_interpolation_matrix`.
+            If provided, this matrix will be applied directly without recomputing
+            the forward or spherical spline mappings. Channel names and bad
+            channels must match the precomputed matrix.
+
+            .. versionadded:: 1.14
         verbose : bool | str | int | None
             Control verbosity of the logging output. If ``None``, use the default
             verbosity level. See the :ref:`logging documentation <tut-logging>` and
@@ -1127,6 +1137,7 @@ def interpolate_bads(
         .. versionadded:: 0.9.0
         """
         from .interpolation import (
+            _do_interp_dots,
             _interpolate_bads_ecog,
             _interpolate_bads_eeg,
             _interpolate_bads_meeg,
@@ -1138,6 +1149,23 @@ def interpolate_bads(
         _check_preload(self, "interpolation")
         _validate_type(method, (dict, str, None), "method")
 
+        if interpolation is not None:
+            _validate_type(interpolation, (dict,), "interpolation")
+            if interpolation.get("ch_names") != list(self.ch_names) or set(
+                interpolation.get("bads", [])
+            ) != set(self.info["bads"]):
+                raise ValueError(
+                    "The precomputed interpolation matrix is not compatible with "
+                    "this instance. Channel names and bad channels must match."
+                )
+            for step in interpolation.get("steps", []):
+                _do_interp_dots(
+                    self, step["matrix"], step["goods_idx"], step["bads_idx"]
+                )
+            if reset_bads is True:
+                self.info["bads"] = [ch for ch in self.info["bads"] if ch in exclude]
+            return self
+
         # check for channels with invalid position(s)
         invalid_chs = []
         for ch in self.info["bads"]:

mne/channels/interpolation.py

@@ -9,12 +9,18 @@
 from scipy.interpolate import RectBivariateSpline
 from scipy.linalg import pinv
 
-from .._fiff.meas_info import _simplify_info, create_info
+from .._fiff.meas_info import Info, _simplify_info, create_info
 from .._fiff.pick import pick_channels, pick_info, pick_types
 from .._fiff.proj import _has_eeg_average_ref_proj, make_eeg_average_ref_proj
 from ..bem import _check_origin
 from ..surface import _normalize_vectors
-from ..utils import _validate_type, _verbose_control, logger, warn
+from ..utils import (
+    _validate_type,
+    _verbose_control,
+    logger,
+    object_diff,
+    warn,
+)
 
 
 def _calc_h(cosang, stiffness=4, n_legendre_terms=50):
@@ -133,26 +139,72 @@ def _do_interp_dots(inst, interpolation, goods_idx, bads_idx):
     )
 
 
-@_verbose_control
-def _interpolate_bads_eeg(inst, origin, exclude=None, ecog=False, verbose=None):
+class InterpolationMatrix(dict):
+    """Channel interpolation matrix and metadata.
+
+    Parameters
+    ----------
+    steps : list of dict
+        The interpolation steps. Each step is a dict containing
+        ``'matrix'``, ``'goods_idx'``, ``'bads_idx'``, ``'ch_type'``, and ``'method'``.
+    ch_names : list of str
+        The channel names of the instance.
+    bads : list of str
+        The bad channels.
+    origin : tuple | ndarray | str | None
+        The origin used for interpolation.
+    """
+
+    def __init__(self, steps, ch_names, bads, origin=None):
+        super().__init__(
+            steps=list(steps),
+            ch_names=list(ch_names),
+            bads=list(bads),
+            origin=origin,
+        )
+
+    def __repr__(self):  # noqa: D105
+        ch_types = ", ".join(
+            t for s in self.get("steps", []) if (t := s.get("ch_type"))
+        )
+        types_info = f" ({ch_types})" if ch_types else ""
+        return (
+            f"<InterpolationMatrix | {len(self.get('bads', []))} bad channels, "
+            f"{len(self.get('steps', []))} steps{types_info}>"
+        )
+
+    def copy(self):
+        """Return a copy of the InterpolationMatrix.
+
+        Returns
+        -------
+        inst : instance of InterpolationMatrix
+            A copy of the interpolation matrix.
+        """
+        return deepcopy(self)
+
+    def __eq__(self, other):  # noqa: D105
+        return isinstance(other, InterpolationMatrix) and object_diff(self, other) == ""
+
+
+def _compute_spline_interpolation(info, origin, exclude=(), ecog=False):
+    """Compute spherical spline interpolation matrix for EEG or ECoG."""
     if exclude is None:
         exclude = list()
-    bads_idx = np.zeros(len(inst.ch_names), dtype=bool)
-    goods_idx = np.zeros(len(inst.ch_names), dtype=bool)
+    bads_idx = np.zeros(len(info["ch_names"]), dtype=bool)
+    goods_idx = np.zeros(len(info["ch_names"]), dtype=bool)
 
-    picks = pick_types(inst.info, meg=False, eeg=not ecog, ecog=ecog, exclude=exclude)
-    inst.info._check_consistency()
-    bads_idx[picks] = [inst.ch_names[ch] in inst.info["bads"] for ch in picks]
+    picks = pick_types(info, meg=False, eeg=not ecog, ecog=ecog, exclude=exclude)
+    info._check_consistency()
+    bads_idx[picks] = [info["ch_names"][ch] in info["bads"] for ch in picks]
 
     if len(picks) == 0 or bads_idx.sum() == 0:
-        return
+        return None
 
     goods_idx[picks] = True
     goods_idx[bads_idx] = False
 
-    pos = inst._get_channel_positions(picks)
-
-    # Make sure only EEG are used
+    pos = info._get_channel_positions(picks)
     bads_idx_pos = bads_idx[picks]
     goods_idx_pos = goods_idx[picks]
 
@@ -169,8 +221,75 @@ def _interpolate_bads_eeg(inst, origin, exclude=None, ecog=False, verbose=None):
     pos_bad = pos[bads_idx_pos] - origin
     logger.info(f"Computing interpolation matrix from {len(pos_good)} sensor positions")
     interpolation = _make_interpolation_matrix(pos_good, pos_bad)
+    return goods_idx, bads_idx, interpolation
 
-    logger.info(f"Interpolating {len(pos_bad)} sensors")
+
+@_verbose_control
+def compute_bad_channel_interpolation_matrix(
+    inst_or_info,
+    *,
+    origin="auto",
+    exclude=(),
+    verbose=None,
+):
+    """Compute the bad channel interpolation matrix for EEG.
+
+    Parameters
+    ----------
+    inst_or_info : instance of Raw | Epochs | Evoked | Info
+        The data instance or Info containing channel positions and bad channels.
+    origin : array-like, shape (3,) | str
+        Origin of the sphere in the head coordinate frame and in meters.
+        Can be ``'auto'`` (default), which means a head-digitization-based
+        origin fit.
+    exclude : list | tuple
+        The channels to exclude from interpolation.
+    verbose : bool | str | int | None
+        Control verbosity of the logging output. If ``None``, use the default
+        verbosity level. See the :ref:`logging documentation <tut-logging>` and
+        :func:`mne.verbose` for details. Should only be passed as a keyword
+        argument.
+
+    Returns
+    -------
+    interpolation : instance of InterpolationMatrix
+        The interpolation matrix object that can be passed to
+        :meth:`~mne.io.Raw.interpolate_bads`, :meth:`~mne.Epochs.interpolate_bads`,
+        or :meth:`~mne.Evoked.interpolate_bads`.
+
+    Notes
+    -----
+    .. versionadded:: 1.14
+    """
+    info = getattr(inst_or_info, "info", inst_or_info)
+    _validate_type(info, (Info,), "inst_or_info", "Raw, Epochs, Evoked, or Info")
+
+    resolved_origin = _check_origin(origin, info)
+    res = _compute_spline_interpolation(info, resolved_origin, exclude=exclude)
+    if res is None:
+        warn("No bad channels to interpolate. Doing nothing...")
+        return InterpolationMatrix([], info["ch_names"], info["bads"], origin=None)
+
+    goods_idx, bads_idx, interp_mat = res
+    step = dict(
+        ch_type="eeg",
+        method="spline",
+        goods_idx=goods_idx,
+        bads_idx=bads_idx,
+        matrix=interp_mat,
+    )
+    return InterpolationMatrix(
+        [step], info["ch_names"], info["bads"], origin=resolved_origin
+    )
+
+
+@_verbose_control
+def _interpolate_bads_eeg(inst, origin, exclude=None, ecog=False, verbose=None):
+    res = _compute_spline_interpolation(inst.info, origin, exclude=exclude, ecog=ecog)
+    if res is None:
+        return
+    goods_idx, bads_idx, interpolation = res
+    logger.info(f"Interpolating {bads_idx.sum()} sensors")
     _do_interp_dots(inst, interpolation, goods_idx, bads_idx)
 
 

mne/channels/tests/test_interpolation.py

@@ -684,3 +684,52 @@ def test_interpolate_to_meg(monkeypatch):
         evoked_rt.interpolate_to("ctf151", **kwargs)
     with pytest.raises(ValueError, match="Invalid value"):
         evoked_rt.interpolate_to("foo", **kwargs)
+
+
+def test_cached_interpolation_matrix():
+    """Test caching bad channel interpolation matrix for EEG (gh-4651)."""
+    from mne.channels import (
+        InterpolationMatrix,
+        compute_bad_channel_interpolation_matrix,
+    )
+
+    # Setup raw EEG data with a bad channel
+    raw = read_raw_fif(raw_fname).crop(0, 0.5).load_data()
+    raw.pick(pick_types(raw.info, meg=False, eeg=True)[:10])
+    raw.info["bads"] = [raw.ch_names[2]]
+
+    interp = compute_bad_channel_interpolation_matrix(raw)
+    assert isinstance(interp, InterpolationMatrix)
+    assert repr(interp) == "<InterpolationMatrix | 1 bad channels, 1 steps (eeg)>"
+    assert interp.copy() == interp
+    assert compute_bad_channel_interpolation_matrix(raw.info) == interp
+
+    # Numerical parity across Raw, Epochs, Evoked
+    epochs = make_fixed_length_epochs(raw, duration=0.2, preload=True)
+    for inst in (raw, epochs, epochs.average()):
+        assert_allclose(
+            inst.copy().interpolate_bads(interpolation=interp).get_data(),
+            inst.copy().interpolate_bads().get_data(),
+            atol=1e-15,
+        )
+
+    # Input validation and compatibility checks
+    with pytest.raises(TypeError, match="interpolation must be an instance of"):
+        raw.interpolate_bads(interpolation="invalid")
+    with pytest.raises(TypeError, match="inst_or_info must be an instance of"):
+        compute_bad_channel_interpolation_matrix("invalid")
+    raw_bad = raw.copy()
+    raw_bad.info["bads"] = [raw.ch_names[0]]
+    for bad_inst in (raw.copy().drop_channels([raw.ch_names[0]]), raw_bad):
+        with pytest.raises(ValueError, match="not compatible"):
+            bad_inst.interpolate_bads(interpolation=interp)
+
+    # Empty bads and reset_bads=False
+    info_no_bads = raw.info.copy()
+    info_no_bads["bads"] = []
+    with pytest.warns(RuntimeWarning, match="No bad channels to interpolate"):
+        assert len(compute_bad_channel_interpolation_matrix(info_no_bads)["steps"]) == 0
+    assert (
+        raw.copy().interpolate_bads(interpolation=interp, reset_bads=False).info["bads"]
+        == raw.info["bads"]
+    )