diff --git a/doc/changes/dev/14003.bugfix.rst b/doc/changes/dev/14003.bugfix.rst new file mode 100644 index 00000000000..8a0fa1cf79e --- /dev/null +++ b/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 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`. diff --git a/mne/time_frequency/psd.py b/mne/time_frequency/psd.py index 01d932699a1..7e5f179385a 100644 --- a/mne/time_frequency/psd.py +++ b/mne/time_frequency/psd.py @@ -2,7 +2,6 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. -import warnings from functools import partial import numpy as np @@ -10,7 +9,7 @@ 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,63 @@ 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 - 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 - ] + 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). Preempt 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``. + # A named/tuple window (e.g. "hamming") is regenerated by SciPy at the + # shrunk nperseg, so short spans are handled transparently. An explicit + # ndarray/list window has a fixed length and cannot be shortened to match; + # rather than let SciPy raise a cryptic length-mismatch error, bail clearly. + window_is_array = not isinstance(window, (str, tuple)) + funcs = [] + weights = [] + n_short = 0 + for span in x_splits: + span_len = span.shape[-1] + if span_len < n_per_seg: + if window_is_array: + raise ValueError( + f"A good data span ({span_len} samples) is shorter than " + f"n_per_seg ({n_per_seg}), but `window` is a fixed-length " + "array that cannot be shortened to match it. Pass a shorter " + "`window` array, or reduce n_per_seg (or n_fft) so that every " + "good data span is at least n_per_seg samples long." + ) + 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." + ) 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." - ) - - 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. @@ -300,10 +325,10 @@ def func(*args, **kwargs): ) 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),) diff --git a/mne/time_frequency/tests/test_psd.py b/mne/time_frequency/tests/test_psd.py index 9718f80b153..0e38b03201e 100644 --- a/mne/time_frequency/tests/test_psd.py +++ b/mne/time_frequency/tests/test_psd.py @@ -58,6 +58,111 @@ 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 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) + 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_recovers_band_power(): + """Short good spans must not bias the recovered band power (gh-13039). + + A stationary signal with known 10 Hz power is broken into good spans that are + all shorter than ``n_per_seg``. Shrinking the window per span keeps the alpha + band power close to the value estimated from the uninterrupted signal. + + This asserts on the recovered *power* rather than only on the output being + finite, because the two are not equivalent: zero-padding each short span up to + ``n_per_seg`` also produces finite, plausible-looking output while + underestimating band power badly (the padded zeros enter the window's energy + normalisation), which a finiteness check cannot detect. + """ + sfreq, freq = 100.0, 10.0 + n_fft, n_overlap = 256, 128 + rng = np.random.default_rng(0) + + def signal(n_times): + times = np.arange(n_times) / sfreq + return (np.sin(2 * np.pi * freq * times) + 0.1 * rng.standard_normal(n_times))[ + None, : + ] + + def alpha_power(psds, freqs): + band = (freqs >= 8.0) & (freqs <= 12.0) + return np.trapezoid(psds[0][band], freqs[band]) + + psds, freqs = psd_array_welch( + signal(60 * int(sfreq)), sfreq=sfreq, n_fft=n_fft, n_overlap=n_overlap + ) + reference = alpha_power(psds, freqs) + + # The same signal, chopped into 100-sample good spans (all < n_per_seg). + nan_col = np.full((1, 1), np.nan) + spans = [] + for _ in range(60): + spans += [signal(100), nan_col] + fragmented = np.concatenate(spans[:-1], axis=-1) + + with pytest.warns(RuntimeWarning, match="shorter than n_per_seg"): + psds, freqs = psd_array_welch( + fragmented, sfreq=sfreq, n_fft=n_fft, n_overlap=n_overlap + ) + + assert_allclose(alpha_power(psds, freqs), reference, rtol=0.1) + + +def test_psd_welch_short_span_array_window_raises(): + """A fixed-length ndarray window can't shrink for a 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) + # Named windows are regenerated by SciPy at the shrunk length, but an explicit + # window array is fixed-length and cannot be shortened to fit the 100-sample + # span, so we raise a clear error instead of a cryptic SciPy length mismatch. + with pytest.raises(ValueError, match="fixed-length array"): + psd_array_welch( + x, sfreq=100, n_fft=n_fft, n_overlap=n_overlap, window=np.hamming(n_fft) + ) + # A full-length array window on all-long spans still works (guard is scoped). + 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)