From 5d5c0ed48bdc13e12d1dff402db9abf5ddfbd6e7 Mon Sep 17 00:00:00 2001 From: Ishaan Maheshwari Date: Sat, 24 Oct 2026 05:37:41 +0530 Subject: [PATCH] ENH: Cache bad channel interpolation matrix for EEG (Part 1) (#4651) --- doc/api/preprocessing.rst | 2 + doc/changes/dev/14455.newfeature.rst | 1 + doc/conf.py | 1 + mne/channels/__init__.pyi | 6 + mne/channels/channels.py | 28 +++++ mne/channels/interpolation.py | 147 ++++++++++++++++++++--- mne/channels/tests/test_interpolation.py | 49 ++++++++ 7 files changed, 220 insertions(+), 14 deletions(-) create mode 100644 doc/changes/dev/14455.newfeature.rst diff --git a/doc/api/preprocessing.rst b/doc/api/preprocessing.rst index fdb84530052..23e3e49dd7f 100644 --- a/doc/api/preprocessing.rst +++ b/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 diff --git a/doc/changes/dev/14455.newfeature.rst b/doc/changes/dev/14455.newfeature.rst new file mode 100644 index 00000000000..f59a9892ec3 --- /dev/null +++ b/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`_. diff --git a/doc/conf.py b/doc/conf.py index e64423e5635..b09b07683d9 100644 --- a/doc/conf.py +++ b/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", diff --git a/mne/channels/__init__.pyi b/mne/channels/__init__.pyi index c7aac715b68..c53ac397895 100644 --- a/mne/channels/__init__.pyi +++ b/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, diff --git a/mne/channels/channels.py b/mne/channels/channels.py index 2f70c506dbf..54e7366e66c 100644 --- a/mne/channels/channels.py +++ b/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 ` 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"]: diff --git a/mne/channels/interpolation.py b/mne/channels/interpolation.py index e44742c335b..873c3d0218c 100644 --- a/mne/channels/interpolation.py +++ b/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"" + ) + + 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 ` 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) diff --git a/mne/channels/tests/test_interpolation.py b/mne/channels/tests/test_interpolation.py index a66f9b6f43f..361c6f61148 100644 --- a/mne/channels/tests/test_interpolation.py +++ b/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) == "" + 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"] + )