Repository: mne-tools/mne-python

Issue:
Error in raw.plot_psd() with bad segments
I think #6748 broke PSD plotting when bad segments are present in annotations. The reason is that `psd_welch` now averages PSDs without taking care of NaNs, which it should because bad annotations are automatically transformed into `np.nan`.

```
import mne
from mne.io import read_raw_fif
from mne.datasets import sample
from mne.time_frequency import psd_welch


path = sample.data_path()
raw = read_raw_fif(path + "/MEG/sample/sample_audvis_filt-0-40_raw.fif")
picks = mne.pick_types(raw.info, meg=False, eeg=True)
raw.set_annotations(mne.Annotations(0, 10, "bad_segment"))
psds, freqs = psd_welch(raw, picks=picks)
```

In this example, all elements in `psds` are `np.nan`. Previously, when only mean was supported, we used `np.nanmean` to ignore NaNs. Not sure if there is an easy fix because `scipy.signal.welch` does not support NaNs AFAIK.


Discussion:
Sounds like a regression that needs a tiny test.

IIRC internally `welch` in SciPy just uses the STFT function. So we should just compute the STFT then do our own combining. We can compare the non-NaN version of our output with SciPy's. @cbrnr you up for this?
> Sounds like a regression that needs a tiny test.

Definitely.

> IIRC internally welch in SciPy just uses the STFT function. So we should just compute the STFT then do our own combining. We can compare the non-NaN version of our output with SciPy's. @cbrnr you up for this?

Isn't that basically what we did before? We had only one option then (`np.nanmean`), but it should be pretty easy to implement the second option with `np.nanmedian`. So basically I'd revert most of the changes and just add the median option, right?
> Isn't that basically what we did before? 

I don't remember...

> I'd revert most of the changes and just add the median option, right?

I wouldn't automatically go with the revert route as it did add some functionality we want IIRC, such as new mode options and the ability to pass a callable, etc. (I think?). And there might be a cleaner way to do this than what we had before. But it's possible a revert would be cleanest, not sure.

Basically I would take the code as it is now work the way it should for data with NaN values as much as possible. Whichever way fixes NaN problems while keeping the newer flexibility the simplest way possible is fine by me.

We do have `mne.time_frequency.stft`, so I would personally try to use that and go from there.
Previously, we used `scipy.signal.spectrogram` and then averaged the spectrograms with `np.nanmean`. Now we use `scipy.signal.welch` for both mean and median averaging (which doesn't support NaNs), but added an option to return the unaveraged spectrograms. IMO the simplest fix is to get rid of `scipy.signal.welch` again and just use `scipy.signal.spectrogram` (and implement or own averaging like before, but add `np.nanmedian`).

Not sure why we have `mne.time_frequency.stft` since this is also available in SciPy (we even have our own `fft` function in `fixes` so there might have been problems with the SciPy implementations at one point - maybe worth checking if this is still necessary).
STFT is from 0.19 and we only require 0.18

Agreed `scipy.signal.spectrogram` is a good option
> STFT is from 0.19 and we only require 0.18

OK. I'll make a PR then using `scipy.signal.spectrogram`.

Resolving patch:
diff --git a/mne/time_frequency/psd.py b/mne/time_frequency/psd.py
--- a/mne/time_frequency/psd.py
+++ b/mne/time_frequency/psd.py
@@ -6,28 +6,22 @@
 
 from ..parallel import parallel_func
 from ..io.pick import _picks_to_idx
-from ..utils import logger, verbose, _time_mask
+from ..utils import logger, verbose, _time_mask, _check_option
 from .multitaper import psd_array_multitaper
 
 
-def _spect_func(epoch, n_overlap, n_per_seg, nfft, fs, freq_mask, func):
+def _spect_func(epoch, n_overlap, n_per_seg, nfft, fs, freq_mask, func,
+                average):
     """Aux function."""
     _, _, spect = func(epoch, fs=fs, nperseg=n_per_seg, noverlap=n_overlap,
                        nfft=nfft, window='hamming')
-    return spect[..., freq_mask, :]
-
-
-def _welch_func(epoch, n_overlap, n_per_seg, nfft, fs, freq_mask, average,
-                func):
-    """Aux function."""
-    kws = dict(fs=fs, nperseg=n_per_seg, noverlap=n_overlap, nfft=nfft,
-               window='hamming', average=average)
-
-    if average == 'mean':  # Compatibility with SciPy <1.2
-        del kws['average']
-
-    _, psd = func(epoch, **kws)
-    return psd[..., freq_mask]
+    spect = spect[..., freq_mask, :]
+    # Do the averaging here (per epoch) to save memory
+    if average == 'mean':
+        spect = np.nanmean(spect, axis=-1)
+    elif average == 'median':
+        spect = np.nanmedian(spect, axis=-1)
+    return spect
 
 
 def _check_nfft(n, n_fft, n_per_seg, n_overlap):
@@ -124,6 +118,8 @@ def psd_array_welch(x, sfreq, fmin=0, fmax=np.inf, n_fft=256, n_overlap=0,
     -----
     .. versionadded:: 0.14.0
     """
+    _check_option('average', average, (None, 'mean', 'median'))
+
     dshape = x.shape[:-1]
     n_times = x.shape[-1]
     x = x.reshape(-1, n_times)
@@ -140,34 +136,19 @@ def psd_array_welch(x, sfreq, fmin=0, fmax=np.inf, n_fft=256, n_overlap=0,
     # Parallelize across first N-1 dimensions
     x_splits = np.array_split(x, n_jobs)
 
-    if average in ['mean', 'median']:
-        from scipy.signal import welch
-        parallel, my_welch_func, n_jobs = parallel_func(_welch_func,
-                                                        n_jobs=n_jobs)
-
-        psds = parallel(my_welch_func(d, fs=sfreq, freq_mask=freq_mask,
-                                      n_per_seg=n_per_seg, n_overlap=n_overlap,
-                                      nfft=n_fft, average=average, func=welch)
-                        for d in x_splits)
-        psds = np.concatenate(psds, axis=0)
-        psds.shape = dshape + (-1,)
-    elif average is None:
-        from scipy.signal import spectrogram
-        parallel, my_spect_func, n_jobs = parallel_func(_spect_func,
-                                                        n_jobs=n_jobs)
-
-        f_spect = parallel(my_spect_func(d, n_overlap=n_overlap,
-                                         nfft=n_fft,
-                                         fs=sfreq, freq_mask=freq_mask,
-                                         func=spectrogram,
-                                         n_per_seg=n_per_seg)
-                           for d in x_splits)
-        psds = np.concatenate(f_spect, axis=0)
-        psds.shape = dshape + (len(freqs), -1)
-    else:
-        raise ValueError('average must be one of `mean`, `median`, or None, '
-                         'got {}'.format(average))
-
+    from scipy.signal import spectrogram
+    parallel, my_spect_func, n_jobs = parallel_func(_spect_func, n_jobs=n_jobs)
+
+    f_spect = parallel(my_spect_func(d, n_overlap=n_overlap,
+                                     n_per_seg=n_per_seg, nfft=n_fft, fs=sfreq,
+                                     freq_mask=freq_mask, func=spectrogram,
+                                     average=average)
+                       for d in x_splits)
+    psds = np.concatenate(f_spect, axis=0)
+    shape = dshape + (len(freqs),)
+    if average is None:
+        shape = shape + (-1,)
+    psds.shape = shape
     return psds, freqs
 
 


Test patch:
diff --git a/mne/time_frequency/tests/test_psd.py b/mne/time_frequency/tests/test_psd.py
--- a/mne/time_frequency/tests/test_psd.py
+++ b/mne/time_frequency/tests/test_psd.py
@@ -6,7 +6,7 @@
 
 from mne import pick_types, Epochs, read_events
 from mne.io import RawArray, read_raw_fif
-from mne.utils import run_tests_if_main, requires_version
+from mne.utils import run_tests_if_main
 from mne.time_frequency import psd_welch, psd_multitaper, psd_array_welch
 
 base_dir = op.join(op.dirname(__file__), '..', '..', 'io', 'tests', 'data')
@@ -18,11 +18,11 @@ def test_psd_nan():
     """Test handling of NaN in psd_array_welch."""
     n_samples, n_fft, n_overlap = 2048, 1024, 512
     x = np.random.RandomState(0).randn(1, n_samples)
-    psds, freqs = psd_array_welch(
-        x[:n_fft + n_overlap], float(n_fft), n_fft=n_fft, n_overlap=n_overlap)
-    x[n_fft + n_overlap:] = np.nan  # what Raw.get_data() will give us
-    psds_2, freqs_2 = psd_array_welch(
-        x, float(n_fft), n_fft=n_fft, n_overlap=n_overlap)
+    psds, freqs = psd_array_welch(x[:, :n_fft + n_overlap], float(n_fft),
+                                  n_fft=n_fft, n_overlap=n_overlap)
+    x[:, n_fft + n_overlap:] = np.nan  # what Raw.get_data() will give us
+    psds_2, freqs_2 = psd_array_welch(x, float(n_fft), n_fft=n_fft,
+                                      n_overlap=n_overlap)
     assert_allclose(freqs, freqs_2)
     assert_allclose(psds, psds_2)
     # 1-d
@@ -160,7 +160,6 @@ def test_psd():
         assert (psds_ev.shape == (len(kws['picks']), len(freqs)))
 
 
-@requires_version('scipy', '1.2.0')
 @pytest.mark.parametrize('kind', ('raw', 'epochs', 'evoked'))
 def test_psd_welch_average_kwarg(kind):
     """Test `average` kwarg of psd_welch()."""
@@ -216,11 +215,8 @@ def test_psd_welch_average_kwarg(kind):
     assert psds_mean.shape == psds_unagg.shape[:-1]
     assert_allclose(psds_mean, psds_unagg.mean(axis=-1))
 
-    # SciPy's welch() function corrects the median PSD for its bias relative to
-    # the mean.
-    from scipy.signal.spectral import _median_bias
-    median_bias = _median_bias(psds_unagg.shape[-1])
-    assert_allclose(psds_median, np.median(psds_unagg, axis=-1) / median_bias)
+    # Compare with manual median calculation
+    assert_allclose(psds_median, np.median(psds_unagg, axis=-1))
 
 
 @pytest.mark.slowtest