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 diff --git a/mne/filter.py b/mne/filter.py index 977bc51bf95..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 @@ -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 diff --git a/mne/time_frequency/multitaper.py b/mne/time_frequency/multitaper.py index d917056b2ac..643c6d975ae 100644 --- a/mne/time_frequency/multitaper.py +++ b/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: