Speed up notch filter spectrum fit

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`_.

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
 
@@ -1619,6 +1620,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 +1630,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 +1639,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 +1751,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
 

mne/time_frequency/multitaper.py

@@ -274,12 +274,8 @@ 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 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)
     if n_fft % 2 == 0: