PERF: Custom per-segment FFT for Welch PSD mean averaging

For the common case (average='mean', output='power', no NaNs, data
> 10 MB), bypass scipy.signal.spectrogram entirely and compute the
PSD via batched per-segment rfft. This loops over segments (typically
1-3 for epoched data) instead of rows (potentially 100K+), with a
single batched scipy.fft.rfft call per segment.

On 320 epochs x 376 channels, psd_array_welch goes from ~170ms
(chunked spectrogram, previous commit) to ~116ms (1.5x additional).
Combined with the previous commit: 5047ms -> 116ms (43x total).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: sharifhsn

mne/time_frequency/psd.py

@@ -6,7 +6,8 @@
 from functools import partial
 
 import numpy as np
-from scipy.signal import spectrogram
+from scipy.fft import rfft
+from scipy.signal import get_window, spectrogram
 
 from ..fixes import _reshape_view
 from ..parallel import parallel_func
@@ -60,17 +61,51 @@ def _decomp_aggregate_mask(epoch, func, average, freq_sl):
     return spect
 
 
+def _welch_mean(epoch, fs, window, nperseg, nfft, noverlap, detrend, freq_sl):
+    """Compute Welch PSD with mean averaging via batched per-segment FFT.
+
+    Processes one segment at a time across all rows simultaneously, avoiding
+    the overhead of calling scipy.signal.spectrogram per-row. Each segment
+    iteration uses a single batched rfft call over all rows in the chunk.
+    """
+    win = get_window(window, nperseg)
+    scale = 1.0 / (fs * np.dot(win, win))
+    step = nperseg - noverlap
+    seg_starts = np.arange(0, epoch.shape[-1] - nperseg + 1, step)
+    n_seg = len(seg_starts)
+    n_freqs = nfft // 2 + 1
+    row_bytes = epoch[0].nbytes
+    chunk_size = max(1, int(10e6 / row_bytes))
+    n_rows = epoch.shape[0]
+    result = np.empty((n_rows, n_freqs))
+    for r0 in range(0, n_rows, chunk_size):
+        r1 = min(r0 + chunk_size, n_rows)
+        psd_acc = np.zeros((r1 - r0, n_freqs))
+        for s in seg_starts:
+            seg = epoch[r0:r1, s : s + nperseg].copy()
+            if detrend:
+                seg -= seg.mean(axis=-1, keepdims=True)
+            seg *= win
+            ft = rfft(seg, n=nfft, axis=-1)
+            psd_acc += np.real(ft * ft.conj())
+        psd_acc *= scale / n_seg
+        # One-sided spectrum: double non-DC/Nyquist components
+        if nfft % 2:
+            psd_acc[:, 1:] *= 2
+        else:
+            psd_acc[:, 1:-1] *= 2
+        result[r0:r1] = psd_acc
+    return result[:, freq_sl]
+
+
 def _spect_func(epoch, func, freq_sl, average, *, output="power"):
     """Aux function."""
-    # Process in chunks to balance vectorization (scipy.signal.spectrogram
-    # handles multi-row input efficiently) against memory usage.
     kwargs = dict(func=func, average=average, freq_sl=freq_sl)
     if epoch.nbytes > 10e6:
         # Process in chunks of rows instead of one-by-one. Each chunk is
         # passed to spectrogram as a 2D array, which is much faster than
         # calling spectrogram per-row via np.apply_along_axis.
         n_rows = epoch.shape[0]
-        # Target ~10 MB per chunk (same threshold as the original code)
         row_bytes = epoch[0].nbytes
         chunk_size = max(1, int(10e6 / row_bytes))
         parts = []
@@ -252,69 +287,82 @@ def psd_array_welch(
         f"{n_overlap} overlap and {window} window"
     )
 
-    parallel, my_spect_func, n_jobs = parallel_func(_spect_func, n_jobs=n_jobs)
-    _func = partial(
-        spectrogram,
-        detrend=detrend,
-        noverlap=n_overlap,
-        nperseg=n_per_seg,
-        nfft=n_fft,
-        fs=sfreq,
-        window=window,
-        mode=mode,
-    )
-    if nan_present and aligned_nan:
-        # Aligned NaNs across channels → treat as bad annotations.
-        good_mask = ~nan_mask_full
-        t_onsets, t_offsets = _mask_to_onsets_offsets(good_mask[0])
-        x_splits = [x[..., t_ons:t_off] for t_ons, t_off in zip(t_onsets, t_offsets)]
-        # weights reflect the number of samples used from each span. For spans longer
-        # than `n_per_seg`, trailing samples may be discarded. For spans shorter than
-        # `n_per_seg`, the wrapped function (`scipy.signal.spectrogram`) automatically
-        # reduces `n_per_seg` to match the span length (with a warning).
-        step = n_per_seg - n_overlap
-        span_lengths = [span.shape[-1] for span in x_splits]
-        weights = [
-            w if w < n_per_seg else w - ((w - n_overlap) % step) for w in span_lengths
-        ]
-        agg_func = partial(np.average, weights=weights)
-        if n_jobs > 1:
-            logger.info(
-                f"Data split into {len(x_splits)} (probably unequal) chunks due to "
-                '"bad_*" annotations. Parallelization may be sub-optimal.'
-            )
-        if (np.array(span_lengths) < n_per_seg).any():
-            logger.info(
-                "At least one good data span is shorter than n_per_seg, and will be "
-                "analyzed with a shorter window than the rest of the file."
-            )
+    # Fast path: for the common case of mean-averaged PSD without NaN
+    # complications, compute directly via batched per-segment FFT. This
+    # loops over segments (typically 1-3) rather than rows (potentially
+    # 100K+), using a single batched rfft call per segment.
+    psds = None
+    if average == "mean" and output == "power" and not nan_present and x.nbytes > 10e6:
+        psds = _welch_mean(
+            x, sfreq, window, n_per_seg, n_fft, n_overlap, detrend, freq_sl
+        )
 
-        def func(*args, **kwargs):
-            # swallow SciPy warnings caused by short good data spans
-            with warnings.catch_warnings():
-                warnings.filterwarnings(
-                    action="ignore",
-                    module="scipy",
-                    category=UserWarning,
-                    message=r"nperseg = \d+ is greater than input length",
+    if psds is None:
+        parallel, my_spect_func, n_jobs = parallel_func(_spect_func, n_jobs=n_jobs)
+        _func = partial(
+            spectrogram,
+            detrend=detrend,
+            noverlap=n_overlap,
+            nperseg=n_per_seg,
+            nfft=n_fft,
+            fs=sfreq,
+            window=window,
+            mode=mode,
+        )
+        if nan_present and aligned_nan:
+            # Aligned NaNs across channels → treat as bad annotations.
+            good_mask = ~nan_mask_full
+            t_onsets, t_offsets = _mask_to_onsets_offsets(good_mask[0])
+            x_splits = [
+                x[..., t_ons:t_off] for t_ons, t_off in zip(t_onsets, t_offsets)
+            ]
+            # weights reflect the number of samples used from each span.
+            step = n_per_seg - n_overlap
+            span_lengths = [span.shape[-1] for span in x_splits]
+            weights = [
+                w if w < n_per_seg else w - ((w - n_overlap) % step)
+                for w in span_lengths
+            ]
+            agg_func = partial(np.average, weights=weights)
+            if n_jobs > 1:
+                logger.info(
+                    f"Data split into {len(x_splits)} (probably unequal) chunks "
+                    'due to "bad_*" annotations. Parallelization may be '
+                    "sub-optimal."
+                )
+            if (np.array(span_lengths) < n_per_seg).any():
+                logger.info(
+                    "At least one good data span is shorter than n_per_seg, "
+                    "and will be analyzed with a shorter window than the rest "
+                    "of the file."
                 )
-                return _func(*args, **kwargs)
 
-    else:
-        # Either no NaNs, or NaNs are not aligned across channels.
-        if nan_present and not aligned_nan:
-            logger.info(
-                "NaN masks are not aligned across channels; treating NaNs as "
-                "per-channel contamination."
-            )
-        x_splits = [arr for arr in np.array_split(x, n_jobs) if arr.size != 0]
-        agg_func = np.concatenate
-        func = _func
-    f_spect = parallel(
-        my_spect_func(d, func=func, freq_sl=freq_sl, average=average, output=output)
-        for d in x_splits
-    )
-    psds = agg_func(f_spect, axis=0)
+            def func(*args, **kwargs):
+                # swallow SciPy warnings caused by short good data spans
+                with warnings.catch_warnings():
+                    warnings.filterwarnings(
+                        action="ignore",
+                        module="scipy",
+                        category=UserWarning,
+                        message=r"nperseg = \d+ is greater than input length",
+                    )
+                    return _func(*args, **kwargs)
+
+        else:
+            # Either no NaNs, or NaNs are not aligned across channels.
+            if nan_present and not aligned_nan:
+                logger.info(
+                    "NaN masks are not aligned across channels; treating NaNs "
+                    "as per-channel contamination."
+                )
+            x_splits = [arr for arr in np.array_split(x, n_jobs) if arr.size != 0]
+            agg_func = np.concatenate
+            func = _func
+        f_spect = parallel(
+            my_spect_func(d, func=func, freq_sl=freq_sl, average=average, output=output)
+            for d in x_splits
+        )
+        psds = agg_func(f_spect, axis=0)
     shape = dshape + (len(freqs),)
     if average is None:
         shape = shape + (-1,)