Fix psd_array_welch for good data spans shorter than n_per_seg (#13039)

doc/changes/dev/14003.bugfix.rst

@@ -0,0 +1 @@
+Fix :func:`mne.time_frequency.psd_array_welch` (and Welch-method ``compute_psd``) so that good data spans shorter than ``n_per_seg`` no longer raise ``noverlap must be less than nperseg``; such spans are now zero-padded up to ``n_per_seg`` (with a warning) instead of being lost, by :newcontrib:`Cedric Conday`.

mne/time_frequency/psd.py

@@ -2,15 +2,14 @@
 # License: BSD-3-Clause
 # Copyright the MNE-Python contributors.
 
-import warnings
 from functools import partial
 
 import numpy as np
 from scipy.signal import spectrogram
 
 from ..fixes import _reshape_view
 from ..parallel import parallel_func
-from ..utils import _check_option, _ensure_int, logger, verbose, warn
+from ..utils import _check_option, _ensure_int, _pl, logger, verbose, warn
 from ..utils.numerics import _mask_to_onsets_offsets
 
 
@@ -259,37 +258,49 @@ def psd_array_welch(
         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
+        if not x_splits:
+            raise ValueError(
+                "No good data spans remain to compute the PSD (all samples are "
+                "excluded by bad annotations)."
+            )
+        # A good data span shorter than n_per_seg cannot hold a full-length Welch
+        # window. SciPy clamps nperseg down to the span length internally but leaves
+        # noverlap unchanged, then raises "noverlap must be less than nperseg"
+        # (#13039). Zero-pad each short span up to n_per_seg instead. This matches
+        # the direction SciPy itself is taking for short input (scipy#25608; scipy
+        # PR #25633 zero-pads the csd input up to nperseg rather than shrinking the
+        # window), so every span keeps the same window length and frequency grid and
+        # an explicit ``window`` array still fits. A padded span yields a single
+        # full-length segment; its weight counts only the real (pre-padding) samples
+        # so the zeros do not inflate its share of the average.
         span_lengths = [span.shape[-1] for span in x_splits]
+        n_short = sum(span_len < n_per_seg for span_len in span_lengths)
+        if n_short:
+            x_splits = [
+                np.pad(span, [(0, 0)] * (span.ndim - 1) + [(0, n_per_seg - span_len)])
+                if span_len < n_per_seg
+                else span
+                for span, span_len in zip(x_splits, span_lengths)
+            ]
+            warn(
+                f"{n_short} good data span{_pl(n_short)} shorter than n_per_seg "
+                f"({n_per_seg}) {'was' if n_short == 1 else 'were'} zero-padded up to "
+                "n_per_seg; the padding biases their spectral estimates. Reduce "
+                "n_per_seg (or n_fft) to avoid this."
+            )
+        step = n_per_seg - n_overlap
+        # weight each span by the samples covered by whole windows (trailing
+        # remainder is discarded); a zero-padded short span counts its real length.
         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)
+        func = _func
         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."
-            )
-
-        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.

mne/time_frequency/tests/test_psd.py

@@ -58,6 +58,66 @@ def test_bad_annot_handling():
     np.testing.assert_allclose(got[0], want[0], rtol=1e-15, atol=0)
 
 
+def test_psd_welch_short_span_kept():
+    """Good spans shorter than n_per_seg are kept with a warning (gh-13039)."""
+    n_fft = 256
+    n_overlap = n_fft // 2  # 128
+    n_chan = 2
+    rng = np.random.default_rng(0)
+    # A short good span (100 samples < n_per_seg), then a bad-annotation
+    # (aligned NaN), then a long span. The short span cannot hold a full Welch
+    # window; instead of raising from SciPy ("noverlap must be less than
+    # nperseg") or being dropped, it is now zero-padded up to n_per_seg (matching
+    # SciPy's own direction for short input, scipy#25608) with a warning.
+    short = rng.standard_normal((n_chan, 100))
+    long = rng.standard_normal((n_chan, 5 * n_fft))
+    x = np.concatenate((short, np.full((n_chan, 1), np.nan), long), axis=-1)
+    with pytest.warns(RuntimeWarning, match="shorter than n_per_seg"):
+        psds, freqs = psd_array_welch(x, sfreq=100, n_fft=n_fft, n_overlap=n_overlap)
+    assert psds.shape == (n_chan, len(freqs))
+    assert np.all(np.isfinite(psds))
+
+    # Even when *every* good span is shorter than n_per_seg, each is analyzed on
+    # its own; the estimate is still computed (previously this raised). Use three
+    # short spans so the total length exceeds n_fft (else the n_fft > n_times
+    # guard fires first).
+    nan_col = np.full((n_chan, 1), np.nan)
+    x_all_short = np.concatenate((short, nan_col, short, nan_col, short), axis=-1)
+    with pytest.warns(RuntimeWarning, match="shorter than n_per_seg"):
+        psds, freqs = psd_array_welch(
+            x_all_short, sfreq=100, n_fft=n_fft, n_overlap=n_overlap
+        )
+    assert psds.shape == (n_chan, len(freqs))
+    assert np.all(np.isfinite(psds))
+
+
+def test_psd_welch_short_span_array_window():
+    """A fixed-length ndarray window still fits a zero-padded short span (gh-13039)."""
+    n_fft = 256
+    n_overlap = n_fft // 2
+    n_chan = 2
+    rng = np.random.default_rng(0)
+    short = rng.standard_normal((n_chan, 100))
+    long = rng.standard_normal((n_chan, 5 * n_fft))
+    x = np.concatenate((short, np.full((n_chan, 1), np.nan), long), axis=-1)
+    # Zero-padding the short span up to n_per_seg means an explicit full-length
+    # window array (which SciPy requires to equal nperseg) now fits it too, so the
+    # earlier fixed-length-window special case is unnecessary; it just warns.
+    with pytest.warns(RuntimeWarning, match="shorter than n_per_seg"):
+        psds, freqs = psd_array_welch(
+            x, sfreq=100, n_fft=n_fft, n_overlap=n_overlap, window=np.hamming(n_fft)
+        )
+    assert psds.shape == (n_chan, len(freqs))
+    assert np.all(np.isfinite(psds))
+    # A full-length array window on all-long spans still works (no padding, no warn).
+    x_ok = np.concatenate((long, np.full((n_chan, 1), np.nan), long), axis=-1)
+    psds, freqs = psd_array_welch(
+        x_ok, sfreq=100, n_fft=n_fft, n_overlap=n_overlap, window=np.hamming(n_fft)
+    )
+    assert psds.shape == (n_chan, len(freqs))
+    assert np.all(np.isfinite(psds))
+
+
 def _make_psd_data():
     """Make noise data with sinusoids in 2 out of 7 channels."""
     rng = np.random.default_rng(0)