Repository: mne-tools/mne-python

Issue:
All NaNs in `raw.plot_psd()`
### Description of the problem

I have a (large) EEG file, which works fine when I compute and plot its PSD with `raw.plot_psd()`. However, I then manually mark some segments as bad, perform ICA, remove some ICs, and back-project. Using this data, I cannot compute the PSD anymore:

```
ValueError                                Traceback (most recent call last)
Cell In[3], line 1
----> 1 raw.plot_psd()

File <decorator-gen-213>:5, in plot_psd(self, fmin, fmax, tmin, tmax, picks, proj, reject_by_annotation, method, average, dB, estimate, xscale, area_mode, area_alpha, color, line_alpha, spatial_colors, sphere, exclude, ax, show, n_jobs, verbose, **method_kw)

File <decorator-gen-212>:12, in plot_psd(self, fmin, fmax, tmin, tmax, picks, proj, reject_by_annotation, method, average, dB, estimate, xscale, area_mode, area_alpha, color, line_alpha, spatial_colors, sphere, exclude, ax, show, n_jobs, verbose, **method_kw)

File ~/Projects/mne-python/mne/time_frequency/spectrum.py:114, in SpectrumMixin.plot_psd(self, fmin, fmax, tmin, tmax, picks, proj, reject_by_annotation, method, average, dB, estimate, xscale, area_mode, area_alpha, color, line_alpha, spatial_colors, sphere, exclude, ax, show, n_jobs, verbose, **method_kw)
    111 ci = 'sd' if area_mode == 'std' else area_mode
    112 # ↓ here picks="all" because we've already restricted the `info` to
    113 # ↓ have only `picks` channels
--> 114 fig = spectrum.plot(
    115     picks='all', average=average, dB=dB, amplitude=amplitude,
    116     xscale=xscale, ci=ci, ci_alpha=area_alpha, color=color,
    117     alpha=line_alpha, spatial_colors=spatial_colors, sphere=sphere,
    118     exclude=exclude, axes=ax, show=show)
    119 return fig

File ~/Projects/mne-python/mne/time_frequency/spectrum.py:611, in BaseSpectrum.plot(self, picks, average, dB, amplitude, xscale, ci, ci_alpha, color, alpha, spatial_colors, sphere, exclude, axes, show)
    609 xlabels_list = [False] * (len(axes) - 1) + [True]
    610 # plot
--> 611 _plot_psd(self, fig, self.freqs, psd_list, picks_list, titles_list,
    612           units_list, scalings_list, axes, make_label, color,
    613           area_mode=ci, area_alpha=ci_alpha, dB=dB, estimate=estimate,
    614           average=average, spatial_colors=spatial_colors,
    615           xscale=xscale, line_alpha=alpha,
    616           sphere=sphere, xlabels_list=xlabels_list)
    617 fig.subplots_adjust(hspace=0.3)
    618 plt_show(show, fig)

File ~/Projects/mne-python/mne/viz/utils.py:2277, in _plot_psd(inst, fig, freqs, psd_list, picks_list, titles_list, units_list, scalings_list, ax_list, make_label, color, area_mode, area_alpha, dB, estimate, average, spatial_colors, xscale, line_alpha, sphere, xlabels_list)
   2275     if not spatial_colors:
   2276         spatial_colors = color
-> 2277     _plot_lines(psd_array, info, picks, fig, ax_list, spatial_colors,
   2278                 unit, units=units, scalings=None, hline=None, gfp=False,
   2279                 types=types, zorder='std', xlim=(freqs[0], freqs[-1]),
   2280                 ylim=None, times=freqs, bad_ch_idx=[], titles=titles,
   2281                 ch_types_used=ch_types_used, selectable=True, psd=True,
   2282                 line_alpha=line_alpha, nave=None, time_unit='ms',
   2283                 sphere=sphere, highlight=None)
   2285 for ii, (ax, xlabel) in enumerate(zip(ax_list, xlabels_list)):
   2286     ax.grid(True, linestyle=':')

File ~/Projects/mne-python/mne/viz/evoked.py:462, in _plot_lines(data, info, picks, fig, axes, spatial_colors, unit, units, scalings, hline, gfp, types, zorder, xlim, ylim, times, bad_ch_idx, titles, ch_types_used, selectable, psd, line_alpha, nave, time_unit, sphere, highlight)
    459 if len(idx) > 0:
    460     # Set amplitude scaling
    461     D = this_scaling * data[idx, :]
--> 462     _check_if_nan(D)
    463     gfp_only = gfp == 'only'
    464     if not gfp_only:

File ~/Projects/mne-python/mne/utils/check.py:571, in _check_if_nan(data, msg)
    569 """Raise if any of the values are NaN."""
    570 if not np.isfinite(data).all():
--> 571     raise ValueError("Some of the values {} are NaN.".format(msg))

ValueError: Some of the values  to be plotted are NaN.
```


### Steps to reproduce

```Python
import mne

raw = mne.io.read_raw_fif("test.fif.gz")
raw.plot_psd()
```

### Link to data

Unfortunately, the problematic file is ~240MB. It is available here:

https://drive.google.com/file/d/1q2B5KsS-TKFlBlULBUqswCeNMdn71a-a/view?usp=sharing

### Expected results

Plotting the PSD works.

### Actual results

Plotting the PSD throws an error.

### Additional information

Note that setting `reject_by_annotation=False` works, but I would like to compute the PSD from good segments only (and therefore use the default `reject_by_annotation=True`).

This is on `main`:

<details>

```
Platform:         macOS-11.7.1-x86_64-i386-64bit
Python:           3.10.9 (main, Dec 15 2022, 18:20:40) [Clang 13.0.0 (clang-1300.0.29.30)]
Executable:       /Users/clemens/.local/pipx/venvs/ipython/bin/python
CPU:              i386: 8 cores
Memory:           Unavailable (requires "psutil" package)
mne:              1.4.dev0
numpy:            1.24.1 {unknown linalg bindings (threadpoolctl module not found: No module named 'threadpoolctl')}
scipy:            1.10.0
matplotlib:       3.6.2 {backend=MacOSX}

sklearn:          Not found
numba:            Not found
nibabel:          Not found
nilearn:          Not found
dipy:             Not found
openmeeg:         Not found
cupy:             Not found
pandas:           1.5.2
pyvista:          Not found
pyvistaqt:        Not found
ipyvtklink:       Not found
vtk:              Not found
qtpy:             Not found
ipympl:           Not found
pyqtgraph:        Not found
pooch:            v1.6.0

mne_bids:         Not found
mne_nirs:         Not found
mne_features:     Not found
mne_qt_browser:   Not found
mne_connectivity: Not found
mne_icalabel:     Not found
```
</details>


Discussion:
I suspect that ica.apply forgets to put back the annotations to the
backprojected instance.

> Message ID: ***@***.***>
>

> I suspect that ica.apply forgets to put back the annotations to the backprojected instance.

They are still here in `raw.annotations` – or did you mean something else?
then I would need a full snippet to see if I can replicate. Maybe
first_samp is off?

> Message ID: ***@***.***>
>

I've included a link to a data set and three lines of code to reproduce at the top. I also checked `raw.first_samp` and `raw.first_time`, they are both `0` (and identical to pre-ICA).

I am now pretty sure it has nothing to do with ICA, because I have a data set that gives the same error *prior* to ICA. I suspect it has to do with creating bad annotations, maybe too many, overlapping, or in the wrong place?
It has to do with the `NaN` introduced by `reject_by_annotation` for your `bad_intertrial` annotations.
The spectrogram is computed per channels with `scipy.signal.spectrogram`, and is returning only NaN in this case.

```
import mne
from scipy.signal import spectrogram

raw = mne.io.read_raw_fif(
    "/Users/scheltie/Downloads/test.fif.gz", preload=True
)
data = raw.get_data(picks=0, reject_by_annotation="NaN")[0, :]  # select first channel
_, _, spect = spectrogram(data, fs=raw.info["sfreq"], nperseg=2048, noverlap=0, nfft=2048)
assert np.all(np.isnan(spect))
```

And for some reason, reducing the window size fixes this problem. 

```
_, _, spect = spectrogram(data, fs=raw.info["sfreq"], nperseg=1024, noverlap=0, nfft=1024)
assert not np.all(np.isnan(spect))
```

For this reason, `raw.compute_psd(method="welch", n_fft=2048).plot()` fails (corresponds to `raw.plot_psd`) while `raw.compute_psd(method="welch", n_fft=1024).plot()` works. Note that the failure occurs in the plotting function, since the Spectrum object is created even if all the values in `spectrum._data` are `NaN`.

That's all the digging I can do for now.. Anyone knows what `scipy.signal.spectrogram` is doing on arrays that contain `NaN` values?
Thanks for investigating @mscheltienne. So it looks like maybe there is a problem when the remaining data (the good segments) are too short?
Meaning they are shorter than `nfft`? Good idea, here is the distribution of the segments length on your file:

![Figure_1](https://user-images.githubusercontent.com/73893616/212327529-9c3db4af-2aed-454c-be05-8bcab7ffbb5c.png)

They are all shorter than `2048`, maybe that's why? But I don't understand why the very short segments, e.g. `< 1024` do not emit a warning for mean of empty slice with `nfft` set to `1024`.
@cbrnr you have mostly bad segments in this data. What @mscheltienne suggests is correct to me.

this works:

`raw.plot_psd(n_fft=512)`


Yes, this makes sense. Should we add a check to raise a better error? I'd have to do some more testing to find out exactly how to trigger it, but it might be worth it.
Yes, it would be good to raise a better error when plotting an array of `NaN`, and also when creating a `Spectrum` object with an array containing only `NaN`.
Here's a minimal example so we don't have to use my data set:

```python
import mne
from mne.datasets import sample

fname = sample.data_path() / "MEG" / "sample" / "sample_audvis_raw.fif"
raw = mne.io.read_raw(fname).pick_types(eeg=True).crop(0, 10).load_data()

raw.plot_psd()  # works
raw.set_annotations(mne.Annotations(2, 2, "bad"))
raw.plot_psd()
```

This is 10s of data with a sampling frequency of approximately 600.615Hz. Plotting the PSD without annotations works, but it doesn't when there is one bad annotation starting at second 2 with a duration of 2 seconds. That's a bit weird.
I can't replicate getting all NaN / empty plot, using that code snippet on latest `main` I get this plot:

<img width="1112" alt="Screenshot 2023-01-16 at 6 15 43 AM" src="https://user-images.githubusercontent.com/2365790/212698943-c3d9d231-51fe-4740-b095-148a5642734e.png">

... oh wait I didn't notice that there was actually a traceback!
```
/Users/larsoner/python/mne-python/mne/time_frequency/psd.py:51: RuntimeWarning: Mean of empty slice
  spect = np.nanmean(spect, axis=-1)
Traceback (most recent call last):
  File "/Users/larsoner/Desktop/rep.py", line 9, in <module>
    raw.plot_psd()
  File "<decorator-gen-213>", line 5, in plot_psd
  File "<decorator-gen-212>", line 12, in plot_psd
  File "/Users/larsoner/python/mne-python/mne/time_frequency/spectrum.py", line 114, in plot_psd
    fig = spectrum.plot(
  File "/Users/larsoner/python/mne-python/mne/time_frequency/spectrum.py", line 611, in plot
    _plot_psd(self, fig, self.freqs, psd_list, picks_list, titles_list,
  File "/Users/larsoner/python/mne-python/mne/viz/utils.py", line 2278, in _plot_psd
    _plot_lines(psd_array, info, picks, fig, ax_list, spatial_colors,
  File "/Users/larsoner/python/mne-python/mne/viz/evoked.py", line 462, in _plot_lines
    _check_if_nan(D)
  File "/Users/larsoner/python/mne-python/mne/utils/check.py", line 571, in _check_if_nan
    raise ValueError("Some of the values {} are NaN.".format(msg))
ValueError: Some of the values  to be plotted are NaN.
```
I'll look into it a bit
Okay it's because:
```
>>> 601 * 2  # bad annotation start sample (ish)
1202
>>> 601 * 4 # bad annotation stop (ish)
2404
>>> 601 * 10  # total data length
6010
>>> 2048 * 3  # data length needed to fit three len-2048 windows
6144
```
And it looks like `n_per_seg=2048, n_overlap=0` by default. So we construct consecutive length-2048 windows, and only 2 fit, both of which overlap with the BAD annotation, which shows up in the data as NaN.

However, this doesn't fit with the description about how bad data are omitted, or with how bad data are typically omitted (e.g., during filtering operations). Instead of just setting the bad sections to NaN, constructing consecutive windows, and discarding any with `NaN` overlap, it would be better (and I think a bugfix) to take each good segment and chop it up independently. This will also be more correct for example when raw data are concatenated -- we shouldn't allow any window to span the concat point(s), and I think it will with our current code. If we made this change, you'd end up with one usable window from sample 2404-ish (first good sample after the BAD annotation at around 4s) to 2048 samples past that (4452), which makes more sense to me. @drammock agreed?

In the meantime @cbrnr you could work around this for example by setting `n_overlap=1024` or otherwise tweaking the overlap and n_fft params.
> And it looks like `n_per_seg=2048`, `n_overlap=0` by default. So we construct consecutive length-2048 windows, and only 2 fit, both of which overlap with the BAD annotation, which shows up in the data as NaN.

OK, I get it. Just to clarify, the last portion of the data (which is shorter than 2048) isn't factoring into the PSD calculation at all?

I agree that chopping up good data segments (maybe up to a specific shortest length) would be much better than the current approach of discarding an entire window if there is just one NaN in it.

Thanks for the workaround, I'll do that in the meantime. It might also be worth discussing the defaults when this gets fixed anyway (maybe overlapping shorter segments would make life a bit easier?).
> Just to clarify, the last portion of the data (which is shorter than 2048) isn't factoring into the PSD calculation at all?-

Yep, which on the one hand is unfortunate because you're almost always going to discard some samples at the end of your data, but on the other hand is consistent with how `n_per_seg` and `n_overlap` are defined and used (adding an extra window to always capture the end samples would violate `n_overlap` in the vast majority of cases)
> it would be better (and I think a bugfix) to take each good segment and chop it up independently.

agreed. I'll add it to my queue but won't likely get to it until next week, feel free @cbrnr or @mscheltienne or @larsoner to tackle it in the meantime if you have bandwidth.

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
@@ -11,6 +11,7 @@
 
 from ..parallel import parallel_func
 from ..utils import _check_option, _ensure_int, logger, verbose
+from ..utils.numerics import _mask_to_onsets_offsets
 
 
 # adapted from SciPy
@@ -224,12 +225,30 @@ def psd_array_welch(
         window=window,
         mode=mode,
     )
-    x_splits = [arr for arr in np.array_split(x, n_jobs) if arr.size != 0]
+    if np.any(np.isnan(x)):
+        good_mask = ~np.isnan(x)
+        # NaNs originate from annot, so must match for all channels. Note that we CANNOT
+        # use np.testing.assert_allclose() here; it is strict about shapes/broadcasting
+        assert np.allclose(good_mask, good_mask[[0]], equal_nan=True)
+        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 = [
+            split.shape[-1] for split in x_splits if split.shape[-1] >= n_per_seg
+        ]
+        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.'
+            )
+    else:
+        x_splits = [arr for arr in np.array_split(x, n_jobs) if arr.size != 0]
+        agg_func = np.concatenate
     f_spect = parallel(
         my_spect_func(d, func=func, freq_sl=freq_sl, average=average, output=output)
         for d in x_splits
     )
-    psds = np.concatenate(f_spect, axis=0)
+    psds = agg_func(f_spect, axis=0)
     shape = dshape + (len(freqs),)
     if average is None:
         shape = shape + (-1,)


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
@@ -36,6 +36,26 @@ def test_psd_nan():
     assert "hamming window" in log
 
 
+def test_bad_annot_handling():
+    """Make sure results equivalent with/without Annotations."""
+    n_per_seg = 256
+    n_chan = 3
+    n_times = 5 * n_per_seg
+    x = np.random.default_rng(seed=42).standard_normal(size=(n_chan, n_times))
+    want = psd_array_welch(x, sfreq=100)
+    # now simulate an annotation that breaks up the array into unequal spans. Using
+    # `n_per_seg` as the cut point is unrealistic/idealized, but it allows us to test
+    # whether we get results ~identical to `want` (which we should in this edge case)
+    x2 = np.concatenate(
+        (x[..., :n_per_seg], np.full((n_chan, 1), np.nan), x[..., n_per_seg:]), axis=-1
+    )
+    got = psd_array_welch(x2, sfreq=100)
+    # freqs should be identical
+    np.testing.assert_array_equal(got[1], want[1])
+    # powers should be very very close
+    np.testing.assert_allclose(got[0], want[0], rtol=1e-15, atol=0)
+
+
 def _make_psd_data():
     """Make noise data with sinusoids in 2 out of 7 channels."""
     rng = np.random.default_rng(0)
diff --git a/mne/time_frequency/tests/test_spectrum.py b/mne/time_frequency/tests/test_spectrum.py
--- a/mne/time_frequency/tests/test_spectrum.py
+++ b/mne/time_frequency/tests/test_spectrum.py
@@ -163,12 +163,13 @@ def test_spectrum_reject_by_annot(raw):
     Cannot use raw_spectrum fixture here because we're testing reject_by_annotation in
     .compute_psd() method.
     """
-    spect_no_annot = raw.compute_psd()
+    kw = dict(n_per_seg=512)  # smaller than shortest good span, to avoid warning
+    spect_no_annot = raw.compute_psd(**kw)
     raw.set_annotations(Annotations([1, 5], [3, 3], ["test", "test"]))
-    spect_benign_annot = raw.compute_psd()
+    spect_benign_annot = raw.compute_psd(**kw)
     raw.annotations.description = np.array(["bad_test", "bad_test"])
-    spect_reject_annot = raw.compute_psd()
-    spect_ignored_annot = raw.compute_psd(reject_by_annotation=False)
+    spect_reject_annot = raw.compute_psd(**kw)
+    spect_ignored_annot = raw.compute_psd(**kw, reject_by_annotation=False)
     # the only one that should be different is `spect_reject_annot`
     assert spect_no_annot == spect_benign_annot
     assert spect_no_annot == spect_ignored_annot