From 7d1444af022d85038065aef0019fdbcc67b2a132 Mon Sep 17 00:00:00 2001 From: Clemens Brunner Date: Thu, 30 Jul 2026 11:54:07 +0200 Subject: [PATCH 1/6] Speed up spectrum_fit --- mne/filter.py | 21 +++++++++--------- mne/time_frequency/multitaper.py | 10 ++++----- mne/time_frequency/tests/test_multitaper.py | 24 +++++++++++++++++++-- 3 files changed, 37 insertions(+), 18 deletions(-) diff --git a/mne/filter.py b/mne/filter.py index 977bc51bf95..ecc1aa6d633 100644 --- a/mne/filter.py +++ b/mne/filter.py @@ -1619,6 +1619,8 @@ def _mt_spectrum_proc( if filter_length is None: filter_length = x.shape[-1] filter_length = min(_to_samples(filter_length, sfreq, "", ""), x.shape[-1]) + pick_mask = np.zeros(len(x), dtype=bool) + pick_mask[picks] = True get_wt = partial( _get_window_thresh, sfreq=sfreq, mt_bandwidth=mt_bandwidth, p_value=p_value ) @@ -1627,7 +1629,7 @@ def _mt_spectrum_proc( if n_jobs == 1: freq_list = list() for ii, x_ in enumerate(x): - if ii in picks: + if pick_mask[ii]: x[ii], f = _mt_spectrum_remove_win( x_, sfreq, line_freqs, notch_widths, window_fun, threshold, get_wt ) @@ -1636,7 +1638,7 @@ def _mt_spectrum_proc( data_new = parallel( p_fun(x_, sfreq, line_freqs, notch_widths, window_fun, threshold, get_wt) for xi, x_ in enumerate(x) - if xi in picks + if pick_mask[xi] ) freq_list = [d[1] for d in data_new] data_new = np.array([d[0] for d in data_new]) @@ -1748,17 +1750,16 @@ def _mt_spectrum_remove( indices = np.unique(np.r_[indices_1, indices_2]) rm_freqs = freqs[indices] - fits = list() - for ind in indices: - c = 2 * A[0, ind] - fit = np.abs(c) * np.cos(freqs[ind] * rads + np.angle(c)) - fits.append(fit) - - if len(fits) == 0: + if len(indices) == 0: datafit = 0.0 else: + c = 2 * A[0, indices] # fitted sinusoids are summed, and subtracted from data - datafit = np.sum(fits, axis=0) + datafit = np.sum( + np.abs(c)[:, np.newaxis] + * np.cos(freqs[indices, np.newaxis] * rads + np.angle(c)[:, np.newaxis]), + axis=0, + ) return x - datafit, rm_freqs diff --git a/mne/time_frequency/multitaper.py b/mne/time_frequency/multitaper.py index d917056b2ac..d2dbc6da52e 100644 --- a/mne/time_frequency/multitaper.py +++ b/mne/time_frequency/multitaper.py @@ -274,12 +274,10 @@ def _mt_spectra(x, dpss, sfreq, n_fft=None, remove_dc=True): # only keep positive frequencies freqs = rfftfreq(n_fft, 1.0 / sfreq) - # The following is equivalent to this, but uses less memory: - # x_mt = fftpack.fft(x[:, np.newaxis, :] * dpss, n=n_fft) - n_tapers = dpss.shape[0] if dpss.ndim > 1 else 1 - x_mt = np.zeros(x.shape[:-1] + (n_tapers, len(freqs)), dtype=np.complex128) - for idx, sig in enumerate(x): - x_mt[idx] = rfft(sig[..., np.newaxis, :] * dpss, n=n_fft) + # Compute the tapered FFTs in one batched call across all leading dims. + # This is faster than looping over signals, at the cost of a larger + # temporary than the old per-signal FFT path. + x_mt = rfft(x[..., np.newaxis, :] * dpss, n=n_fft, axis=-1) # Adjust DC and maybe Nyquist, depending on one-sided transform x_mt[..., 0] /= np.sqrt(2.0) if n_fft % 2 == 0: diff --git a/mne/time_frequency/tests/test_multitaper.py b/mne/time_frequency/tests/test_multitaper.py index 426b148f86a..030d741a92f 100644 --- a/mne/time_frequency/tests/test_multitaper.py +++ b/mne/time_frequency/tests/test_multitaper.py @@ -4,10 +4,11 @@ import numpy as np import pytest -from numpy.testing import assert_array_almost_equal +from numpy.testing import assert_array_almost_equal, assert_allclose +from scipy.fft import rfft from mne.time_frequency import psd_array_multitaper -from mne.time_frequency.multitaper import dpss_windows +from mne.time_frequency.multitaper import _mt_spectra, dpss_windows from mne.utils import _record_warnings @@ -78,3 +79,22 @@ def test_adaptive_weights_convergence(): ): psd_array_multitaper(data, sfreq, adaptive=True, max_iter=2) psd_array_multitaper(data, sfreq, adaptive=True, max_iter=200) + + +def test_mt_spectra_batched_matches_reference(): + """Test that batched tapered spectra match an explicit channel loop.""" + rng = np.random.default_rng(0) + x = rng.standard_normal((3, 17)) + sfreq = 200.0 + dpss, _ = dpss_windows(17, 2.5, 4, low_bias=False) + + x_mt, freqs = _mt_spectra(x, dpss, sfreq) + x = x - np.mean(x, axis=-1, keepdims=True) + ref = np.array([rfft(sig[np.newaxis, :] * dpss, axis=-1) for sig in x]) + ref[..., 0] /= np.sqrt(2.0) + if 17 % 2 == 0: + ref[..., -1] /= np.sqrt(2.0) + + assert x_mt.shape == ref.shape + assert_allclose(x_mt, ref) + assert freqs.shape[0] == x_mt.shape[-1] From 153f18e2ddbf24cbc2420833c6d510ef08bd5c15 Mon Sep 17 00:00:00 2001 From: Clemens Brunner Date: Thu, 30 Jul 2026 12:01:36 +0200 Subject: [PATCH 2/6] Cache _get_window_thresh --- mne/filter.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mne/filter.py b/mne/filter.py index ecc1aa6d633..c3f54f51c7a 100644 --- a/mne/filter.py +++ b/mne/filter.py @@ -6,7 +6,7 @@ from collections import Counter from copy import deepcopy -from functools import partial +from functools import cache, partial from math import gcd import numpy as np @@ -1586,6 +1586,7 @@ def notch_filter( return xf +@cache def _get_window_thresh(n_times, sfreq, mt_bandwidth, p_value): from .time_frequency.multitaper import _compute_mt_params From 0ebb38074c326e548f5de44f70d40deb3f2ad2b1 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:13:43 +0000 Subject: [PATCH 3/6] [autofix.ci] apply automated fixes --- mne/time_frequency/tests/test_multitaper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mne/time_frequency/tests/test_multitaper.py b/mne/time_frequency/tests/test_multitaper.py index 030d741a92f..d1a397e1fdb 100644 --- a/mne/time_frequency/tests/test_multitaper.py +++ b/mne/time_frequency/tests/test_multitaper.py @@ -4,7 +4,7 @@ import numpy as np import pytest -from numpy.testing import assert_array_almost_equal, assert_allclose +from numpy.testing import assert_allclose, assert_array_almost_equal from scipy.fft import rfft from mne.time_frequency import psd_array_multitaper From 4985f17ae790e1b59c46a64e7419e81cf660ed37 Mon Sep 17 00:00:00 2001 From: Clemens Brunner Date: Thu, 30 Jul 2026 12:17:44 +0200 Subject: [PATCH 4/6] Add changelog entry --- doc/changes/dev/14114.bugfix.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 doc/changes/dev/14114.bugfix.rst diff --git a/doc/changes/dev/14114.bugfix.rst b/doc/changes/dev/14114.bugfix.rst new file mode 100644 index 00000000000..165ff641bee --- /dev/null +++ b/doc/changes/dev/14114.bugfix.rst @@ -0,0 +1 @@ +Speed up :func:`mne.filter.notch_filter` and related methods when ``method="spectrum_fit"`` by caching repeated multitaper window setup and batching the tapered FFTs, by `Clemens Brunner`_. \ No newline at end of file From d48fdf67066379b6580034ca65800b3816605016 Mon Sep 17 00:00:00 2001 From: Clemens Brunner Date: Thu, 30 Jul 2026 12:18:56 +0200 Subject: [PATCH 5/6] Remove test (not really needed) --- mne/time_frequency/tests/test_multitaper.py | 24 ++------------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/mne/time_frequency/tests/test_multitaper.py b/mne/time_frequency/tests/test_multitaper.py index d1a397e1fdb..426b148f86a 100644 --- a/mne/time_frequency/tests/test_multitaper.py +++ b/mne/time_frequency/tests/test_multitaper.py @@ -4,11 +4,10 @@ import numpy as np import pytest -from numpy.testing import assert_allclose, assert_array_almost_equal -from scipy.fft import rfft +from numpy.testing import assert_array_almost_equal from mne.time_frequency import psd_array_multitaper -from mne.time_frequency.multitaper import _mt_spectra, dpss_windows +from mne.time_frequency.multitaper import dpss_windows from mne.utils import _record_warnings @@ -79,22 +78,3 @@ def test_adaptive_weights_convergence(): ): psd_array_multitaper(data, sfreq, adaptive=True, max_iter=2) psd_array_multitaper(data, sfreq, adaptive=True, max_iter=200) - - -def test_mt_spectra_batched_matches_reference(): - """Test that batched tapered spectra match an explicit channel loop.""" - rng = np.random.default_rng(0) - x = rng.standard_normal((3, 17)) - sfreq = 200.0 - dpss, _ = dpss_windows(17, 2.5, 4, low_bias=False) - - x_mt, freqs = _mt_spectra(x, dpss, sfreq) - x = x - np.mean(x, axis=-1, keepdims=True) - ref = np.array([rfft(sig[np.newaxis, :] * dpss, axis=-1) for sig in x]) - ref[..., 0] /= np.sqrt(2.0) - if 17 % 2 == 0: - ref[..., -1] /= np.sqrt(2.0) - - assert x_mt.shape == ref.shape - assert_allclose(x_mt, ref) - assert freqs.shape[0] == x_mt.shape[-1] From 9df9b22ded39635c02e5a8fe6ec6f3b1c4ceb4e2 Mon Sep 17 00:00:00 2001 From: Clemens Brunner Date: Thu, 30 Jul 2026 12:22:22 +0200 Subject: [PATCH 6/6] Shorten comment --- mne/time_frequency/multitaper.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/mne/time_frequency/multitaper.py b/mne/time_frequency/multitaper.py index d2dbc6da52e..643c6d975ae 100644 --- a/mne/time_frequency/multitaper.py +++ b/mne/time_frequency/multitaper.py @@ -274,9 +274,7 @@ def _mt_spectra(x, dpss, sfreq, n_fft=None, remove_dc=True): # only keep positive frequencies freqs = rfftfreq(n_fft, 1.0 / sfreq) - # Compute the tapered FFTs in one batched call across all leading dims. - # This is faster than looping over signals, at the cost of a larger - # temporary than the old per-signal FFT path. + # compute FFTs in one batch (faster than looping over tapers but uses more memory) x_mt = rfft(x[..., np.newaxis, :] * dpss, n=n_fft, axis=-1) # Adjust DC and maybe Nyquist, depending on one-sided transform x_mt[..., 0] /= np.sqrt(2.0)