Speed up spectrum_fit

Author: cbrnr

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
 

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:

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]