From 848d1c0197aaa3647a54b53ac4cc0fc123a4f8d2 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 | 41 +++++ mne/channels/interpolation.py | 210 +++++++++++++++++++++-- mne/channels/tests/test_interpolation.py | 60 +++++++ 7 files changed, 307 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..fc642d14f1c 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,36 @@ 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", []): + if step.get("matrix") is None: + self._data[..., step["bads_idx"], :] = np.nan + else: + _do_interp_dots( + self, step["matrix"], step["goods_idx"], step["bads_idx"] + ) + if reset_bads is True: + if any( + step.get("method") == "nan" + for step in interpolation.get("steps", []) + ): + warn( + "interpolate_bads was called with method='nan' and " + "reset_bads=True. Consider setting reset_bads=False so " + "that the nan-containing channels can be easily " + "excluded from later computations." + ) + 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..ef901649a28 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 ( + _on_missing, + _validate_type, + _verbose_control, + logger, + warn, +) def _calc_h(cosang, stiffness=4, n_legendre_terms=50): @@ -133,26 +139,95 @@ 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__() + self["steps"] = list(steps) + self["ch_names"] = list(ch_names) + self["bads"] = list(bads) + self["origin"] = origin + + def __repr__(self): # noqa: D105 + n_steps = len(self.get("steps", [])) + n_bads = len(self.get("bads", [])) + ch_types = [s.get("ch_type", "") for s in self.get("steps", [])] + ch_types_str = ", ".join(t for t in ch_types if t) + types_info = f" ({ch_types_str})" if ch_types_str 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 + if not isinstance(other, InterpolationMatrix): + return False + if ( + self.get("ch_names") != other.get("ch_names") + or self.get("bads") != other.get("bads") + or len(self.get("steps", [])) != len(other.get("steps", [])) + ): + return False + for s1, s2 in zip(self.get("steps", []), other.get("steps", [])): + if s1.get("ch_type") != s2.get("ch_type") or s1.get("method") != s2.get( + "method" + ): + return False + for k in ("goods_idx", "bads_idx", "matrix"): + v1, v2 = s1.get(k), s2.get(k) + if v1 is None or v2 is None: + if v1 is not v2: + return False + elif not np.array_equal(v1, v2): + return False + return True + + def __ne__(self, other): # noqa: D105 + return not 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 +244,115 @@ 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 + + +@_verbose_control +def compute_bad_channel_interpolation_matrix( + inst_or_info, + *, + origin="auto", + method=None, + exclude=(), + on_bad_position="warn", + 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. + method : dict | str | None + Method to use for interpolation. For EEG, supports ``'spline'`` (default). + exclude : list | tuple + The channels to exclude from interpolation. + on_bad_position : "raise" | "warn" | "ignore" + What to do when one or more sensor positions are invalid (zero or NaN). + 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. - logger.info(f"Interpolating {len(pos_bad)} sensors") + 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 + """ + if hasattr(inst_or_info, "info"): + info = inst_or_info.info + else: + info = inst_or_info + _validate_type(info, (Info,), "inst_or_info", "Raw, Epochs, Evoked, or Info") + _validate_type(method, (dict, str, None), "method") + + # Check for channels with invalid position(s) + invalid_chs = [] + for ch in info["bads"]: + if ch in info["ch_names"]: + loc = info["chs"][info["ch_names"].index(ch)]["loc"][:3] + if np.allclose(loc, 0.0, rtol=0, atol=1e-16) or np.isnan(loc).any(): + invalid_chs.append(ch) + + if invalid_chs: + if on_bad_position == "raise": + msg = ( + f"Channel(s) {invalid_chs} have invalid sensor position(s). " + "Interpolation cannot proceed correctly. If you want to " + "continue despite missing positions, set " + "on_bad_position='warn' or 'ignore', which outputs all " + "NaN values (np.nan) for the interpolated channel(s)." + ) + else: + msg = ( + f"Channel(s) {invalid_chs} have invalid sensor position(s) " + "and cannot be interpolated. The values of these channels " + "will be all NaN. To ignore this warning, pass " + "on_bad_position='ignore'." + ) + _on_missing(on_bad_position, msg) + + resolved_origin = _check_origin(origin, info) + res = _compute_spline_interpolation( + info, resolved_origin, exclude=exclude, ecog=False + ) + 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 + steps = [ + dict( + ch_type="eeg", + method="spline", + goods_idx=goods_idx, + bads_idx=bads_idx, + matrix=interp_mat, + ) + ] + return InterpolationMatrix( + steps, 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..2817df4c38a 100644 --- a/mne/channels/tests/test_interpolation.py +++ b/mne/channels/tests/test_interpolation.py @@ -684,3 +684,63 @@ 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, + ) + + # 1. Setup raw EEG data with 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 + + # 2. Test bit-exact numerical parity across Raw, Epochs, Evoked + assert_allclose( + raw.copy().interpolate_bads(interpolation=interp).get_data(), + raw.copy().interpolate_bads().get_data(), + atol=1e-15, + ) + epochs = make_fixed_length_epochs(raw, duration=0.2, preload=True) + assert_allclose( + epochs.copy().interpolate_bads(interpolation=interp).get_data(), + epochs.copy().interpolate_bads().get_data(), + atol=1e-15, + ) + evoked = epochs.average() + assert_allclose( + evoked.copy().interpolate_bads(interpolation=interp).data, + evoked.copy().interpolate_bads().data, + atol=1e-15, + ) + + # 3. Test 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") + with pytest.raises(ValueError, match="not compatible"): + raw.copy().drop_channels([raw.ch_names[0]]).interpolate_bads( + interpolation=interp + ) + raw_bad_mismatch = raw.copy() + raw_bad_mismatch.info["bads"] = [raw.ch_names[0]] + with pytest.raises(ValueError, match="not compatible"): + raw_bad_mismatch.interpolate_bads(interpolation=interp) + + # 4. Empty bads and reset_bads=False + raw_no_bads = raw.copy() + raw_no_bads.info["bads"] = [] + with pytest.warns(RuntimeWarning, match="No bad channels to interpolate"): + assert len(compute_bad_channel_interpolation_matrix(raw_no_bads)["steps"]) == 0 + raw_keep = raw.copy().interpolate_bads(interpolation=interp, reset_bads=False) + assert raw_keep.info["bads"] == raw.info["bads"]