Keep Welch PSD spans shorter than n_per_seg via a per-span window (mne-tools#13039)

Address @drammock's review on mne-tools#14003: rather than dropping good-data spans
shorter than n_per_seg, analyze each such span with nperseg shrunk to the
span length and noverlap clamped below it (a per-span functools.partial
overriding the values baked into _func). n_fft is unchanged, so all spans
share one frequency grid; short spans just get coarser spectral resolution,
which is surfaced via a warning. No data is discarded.

This is the option-3 workaround from the PR discussion. The underlying SciPy
behavior (clamping nperseg but not noverlap) may still be worth reporting
upstream (option 4).

Author: CedricConday

doc/changes/dev/14003.bugfix.rst

@@ -1 +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 dropped from the estimate with a warning, by :newcontrib:`Cedric Conday`.
+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 analyzed with a window shrunk to the span length (with a warning that their spectral resolution is reduced) instead of being lost, by :newcontrib:`Cedric Conday`.

mne/time_frequency/psd.py

@@ -257,37 +257,51 @@ def psd_array_welch(
         # Aligned NaNs across channels → treat as bad annotations.
         good_mask = ~nan_mask_full
         t_onsets, t_offsets = _mask_to_onsets_offsets(good_mask[0])
-        all_splits = [x[..., t_ons:t_off] for t_ons, t_off in zip(t_onsets, t_offsets)]
-        # Drop good data spans shorter than n_per_seg: a single Welch window does not
-        # fit them. (Shrinking the window per-span would mix incompatible estimates,
-        # and passing them to SciPy as-is raises "noverlap must be less than
-        # nperseg".) Warn so the user can lower n_per_seg to keep them. See #13039.
-        x_splits = [span for span in all_splits if span.shape[-1] >= n_per_seg]
-        n_dropped = len(all_splits) - len(x_splits)
-        if n_dropped:
-            warn(
-                f"{n_dropped} good data span{_pl(n_dropped)} shorter than n_per_seg "
-                f"({n_per_seg}) {'was' if n_dropped == 1 else 'were'} excluded from "
-                "the PSD estimate; reduce n_per_seg (or n_fft) to include them."
-            )
+        x_splits = [x[..., t_ons:t_off] for t_ons, t_off in zip(t_onsets, t_offsets)]
         if not x_splits:
             raise ValueError(
-                f"All good data spans are shorter than n_per_seg ({n_per_seg}); no "
-                "data is left to compute the PSD. Reduce n_per_seg (or n_fft)."
+                "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). Pre-empt that here: analyze each short span with a window shrunk
+        # to its own length and noverlap clamped below it. n_fft is left unchanged,
+        # so every span is zero-padded to the same length and shares one frequency
+        # grid; short spans simply get coarser spectral resolution. A per-span
+        # ``partial`` overrides the nperseg/noverlap baked into ``_func``.
+        funcs = []
+        weights = []
+        n_short = 0
+        for span in x_splits:
+            span_len = span.shape[-1]
+            if span_len < n_per_seg:
+                span_nperseg = span_len
+                span_noverlap = min(n_overlap, span_nperseg - 1)
+                n_short += 1
+            else:
+                span_nperseg = n_per_seg
+                span_noverlap = n_overlap
+            funcs.append(partial(_func, nperseg=span_nperseg, noverlap=span_noverlap))
+            # weight by the samples covered by whole windows (trailing remainder is
+            # discarded); a shrunk span contributes a single full-length window.
+            span_step = max(span_nperseg - span_noverlap, 1)
+            weights.append(span_len - ((span_len - span_noverlap) % span_step))
+        if n_short:
+            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'} analyzed with a "
+                "reduced window; spectral resolution is lower for "
+                f"{'that span' if n_short == 1 else 'those spans'}. Reduce n_per_seg "
+                "(or n_fft) to silence this warning."
             )
-        # weights reflect the number of samples used from each (kept) span; trailing
-        # samples beyond the last full window are discarded.
-        step = n_per_seg - n_overlap
-        weights = [
-            w - ((w - n_overlap) % step) for w in (s.shape[-1] for s in x_splits)
-        ]
         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.'
             )
-        func = _func
 
     else:
         # Either no NaNs, or NaNs are not aligned across channels.
@@ -298,10 +312,10 @@ def psd_array_welch(
             )
         x_splits = [arr for arr in np.array_split(x, n_jobs) if arr.size != 0]
         agg_func = np.concatenate
-        func = _func
+        funcs = [_func] * len(x_splits)
     f_spect = parallel(
-        my_spect_func(d, func=func, freq_sl=freq_sl, average=average, output=output)
-        for d in x_splits
+        my_spect_func(d, func=fn, freq_sl=freq_sl, average=average, output=output)
+        for d, fn in zip(x_splits, funcs)
     )
     psds = agg_func(f_spect, axis=0)
     shape = dshape + (len(freqs),)

mne/time_frequency/tests/test_psd.py

@@ -58,16 +58,17 @@ def test_bad_annot_handling():
     np.testing.assert_allclose(got[0], want[0], rtol=1e-15, atol=0)
 
 
-def test_psd_welch_short_span_dropped():
-    """Good spans shorter than n_per_seg are dropped with a warning (gh-13039)."""
+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 single
-    # Welch window; it is now dropped with a warning rather than raising from
-    # SciPy ("noverlap must be less than nperseg").
+    # (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 analyzed with a window shrunk to its
+    # own length and a warning about reduced resolution.
     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)
@@ -76,14 +77,18 @@ def test_psd_welch_short_span_dropped():
     assert psds.shape == (n_chan, len(freqs))
     assert np.all(np.isfinite(psds))
 
-    # If *every* good span is too short, there is nothing left to analyze. Use
-    # three short spans so the total length still exceeds n_fft (otherwise the
-    # earlier n_fft > n_times guard fires first).
+    # 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.raises(ValueError, match="All good data spans are shorter"):
-        with pytest.warns(RuntimeWarning, match="shorter than n_per_seg"):
-            psd_array_welch(x_all_short, sfreq=100, n_fft=n_fft, n_overlap=n_overlap)
+    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 _make_psd_data():