Repository: mne-tools/mne-python

Issue:
mne.time_frequency.read_tfrs() did not work correctly when MNE.Raw.compute_tfr(method='multitaper', output='phase')
### Description of the problem

In the developed version 1.7, after saving the RawTFR object whose output was set as 'phase' with 'save()' method as a h5 file, it cannot be loaded correctly with 'mne.time_frequency.read_tfrs()'.\
The error is 'ValueError: too many values to unpack (expected 3)'\
I found the most possible reason, in Line 4225 of https://github.com/mne-tools/mne-python/blob/main/mne/time_frequency/tfr.py, the data with 4 dimensions were assigned as EpochTFR. But according to https://mne.tools/dev/generated/mne.time_frequency.tfr_array_multitaper.html#mne.time_frequency.tfr_array_multitaper, it says: \
if output in ('complex',' 'phase'), array of shape (n_epochs, n_chans, n_tapers, n_freqs, n_times)\

### Steps to reproduce

```Python
import mne
import numpy as np

# Since my data were also SEEG which were loaded as Raw objects, here I used the SEEG sample for reproduction
misc_path = mne.datasets.misc.data_path()
sample_path = mne.datasets.sample.data_path()
subjects_dir = sample_path / "subjects"
raw = mne.io.read_raw(misc_path / "seeg" / "sample_seeg_ieeg.fif")

# Generate the TFR with output as phase and save the object in the h5 file
spec_session_phase = raw.compute_tfr(method='multitaper', freqs=np.arange(4, 200, 4), output='phase', n_cycles=10, time_bandwidth=4)
spec_session_phase.save("spec_phase.h5", overwrite=True)

# Load the TFR object from the h5 file
spec_session_phase_load = mne.time_frequency.read_tfrs("spec_phase.h5")
```


### Link to data

_No response_

### Expected results

'spec_session_phase_load' should be a RawTFR object as same as spec_session_phase.

### Actual results

ValueError                                Traceback (most recent call last)
Cell In[10], [line 1](vscode-notebook-cell:?execution_count=10&line=1)
----> [1](vscode-notebook-cell:?execution_count=10&line=1) spec_session_phase_load = mne.time_frequency.read_tfrs("spec_phase.h5")

File <decorator-gen-176>:12, in read_tfrs(fname, condition, verbose)

File [/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:4231](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:4231), in read_tfrs(fname, condition, verbose)
   [4229](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:4229) else:
   [4230](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:4230)     Klass = RawTFR
-> [4231](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:4231) out = Klass(inst=hdf5_dict)
   [4232](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:4232) if getattr(out, "metadata", None) is not None:
   [4233](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:4233)     out.metadata = _prepare_read_metadata(out.metadata)

File [/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3213](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3213), in EpochsTFR.__init__(self, info, data, times, freqs, inst, method, comment, tmin, tmax, picks, proj, decim, events, event_id, selection, drop_log, metadata, n_jobs, verbose, **method_kw)
   [3207](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3207) # end TODO ↑↑↑↑↑↑
   [3208](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3208) 
   [3209](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3209) # dict is allowed for __setstate__ compatibility
   [3210](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3210) _validate_type(
   [3211](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3211)     inst, (BaseEpochs, dict), "object passed to EpochsTFR constructor", "Epochs"
   [3212](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3212) )
-> [3213](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3213) super().__init__(
   [3214](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3214)     inst,
   [3215](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3215)     method,
   [3216](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:3216)     freqs,
...
-> [1513](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:1513) n_chan, n_freq, n_time = self._data.shape[self._dims.index("channel") :]
   [1514](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:1514) if n_chan_info != n_chan:
   [1515](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/lib/python3.10/site-packages/mne/time_frequency/tfr.py:1515)     msg = msg.format("Channel", n_chan, "info", n_chan_info)

ValueError: too many values to unpack (expected 3)

### Additional information

Platform             Linux-5.4.0-173-generic-x86_64-with-glibc2.31
Python               3.10.4 (main, Mar 31 2022, 08:41:55) [GCC 7.5.0]
Executable           [/userdata/yhan/myenv/bin/python](https://file+.vscode-resource.vscode-cdn.net/userdata/yhan/myenv/bin/python)
CPU                  x86_64 (48 cores)
Memory               1005.5 GB

Core
├☑ mne               1.7.0.dev163+geee8e6fe5 (devel, latest release is 1.6.1)
├☑ numpy             1.26.4 (OpenBLAS 0.3.23.dev with 48 threads)
├☑ scipy             1.12.0
└☑ matplotlib        3.8.3 (backend=module://matplotlib_inline.backend_inline)

Numerical (optional)
├☑ sklearn           1.4.1.post1
├☑ nibabel           5.2.1
├☑ pandas            2.2.0
├☑ h5io              0.2.2
├☑ h5py              3.10.0
└☐ unavailable       numba, nilearn, dipy, openmeeg, cupy

Visualization (optional)
├☑ pyvista           0.43.3 (OpenGL unavailable)
├☑ pyvistaqt         0.11.0
├☑ vtk               9.3.0
├☑ qtpy              2.4.1 (None=None)
...

Ecosystem (optional)
├☑ mne-bids          0.14
└☐ unavailable       mne-nirs, mne-features, mne-connectivity, mne-icalabel, mne-bids-pipeline, neo

Discussion:
Han-YY: > closes #12527
> 
> @Han-YY hopefully this fixes your issue!

Thanks!

Patch:
diff --git a/doc/changes/devel/12535.bugfix.rst b/doc/changes/devel/12535.bugfix.rst
new file mode 100644
index 000000000..eeeda0bff
--- /dev/null
+++ b/doc/changes/devel/12535.bugfix.rst
@@ -0,0 +1,1 @@
+Calling :meth:`~mne.io.Raw.compute_psd` with ``method="multitaper"`` is now expressly disallowed when ``reject_by_annotation=True`` and ``bad_*`` annotations are present (previously this was nominally allowed but resulted in ``nan`` values in the PSD). By `Daniel McCloy`_.
diff --git a/doc/changes/devel/12536.bugfix.rst b/doc/changes/devel/12536.bugfix.rst
new file mode 100644
index 000000000..2b4a70918
--- /dev/null
+++ b/doc/changes/devel/12536.bugfix.rst
@@ -0,0 +1,1 @@
+:meth:`~mne.io.Raw.compute_psd` and :func:`~mne.time_frequency.psd_array_welch` will now use FFT windows aligned to the onsets of good data spans when ``bad_*`` annotations are present. By `Daniel McCloy`_.
diff --git a/doc/changes/devel/12537.bugfix.rst b/doc/changes/devel/12537.bugfix.rst
new file mode 100644
index 000000000..911bdce44
--- /dev/null
+++ b/doc/changes/devel/12537.bugfix.rst
@@ -0,0 +1,1 @@
+Fix bug in loading of complex/phase TFRs. By `Daniel McCloy`_.
diff --git a/mne/io/base.py b/mne/io/base.py
index c7fb5e4dd..b22b76010 100644
--- a/mne/io/base.py
+++ b/mne/io/base.py
@@ -2197,7 +2197,9 @@ class BaseRaw(
         Parameters
         ----------
         %(method_psd)s
-            Default is ``'welch'``.
+            Note that ``"multitaper"`` cannot be used if ``reject_by_annotation=True``
+            and there are ``"bad_*"`` annotations in the :class:`~mne.io.Raw` data;
+            in such cases use ``"welch"``. Default is ``'welch'``.
         %(fmin_fmax_psd)s
         %(tmin_tmax_psd)s
         %(picks_good_data_noref)s
diff --git a/mne/time_frequency/psd.py b/mne/time_frequency/psd.py
index 33bcd16df..b2123f933 100644
--- a/mne/time_frequency/psd.py
+++ b/mne/time_frequency/psd.py
@@ -11,6 +11,7 @@ from scipy.signal import spectrogram
 
 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,)
diff --git a/mne/time_frequency/spectrum.py b/mne/time_frequency/spectrum.py
index 7300753c5..a9006ac44 100644
--- a/mne/time_frequency/spectrum.py
+++ b/mne/time_frequency/spectrum.py
@@ -1132,6 +1132,12 @@ class Spectrum(BaseSpectrum):
             data = self.inst.get_data(
                 self._picks, start, stop + 1, reject_by_annotation=rba
             )
+            if np.any(np.isnan(data)) and method == "multitaper":
+                raise NotImplementedError(
+                    'Cannot use method="multitaper" when reject_by_annotation=True. '
+                    'Please use method="welch" instead.'
+                )
+
         else:  # Evoked
             data = self.inst.data[self._picks][:, self._time_mask]
         # set nave
diff --git a/mne/time_frequency/tfr.py b/mne/time_frequency/tfr.py
index 00aed68e9..8f8599f75 100644
--- a/mne/time_frequency/tfr.py
+++ b/mne/time_frequency/tfr.py
@@ -1440,7 +1440,7 @@ class BaseTFR(ContainsMixin, UpdateChannelsMixin, SizeMixin, ExtendedTimeMixin):
         self._set_times(self._raw_times)
         # Handle instance type. Prior to gh-11282, Raw was not a possibility so if
         # `inst_type_str` is missing it must be Epochs or Evoked
-        unknown_class = Epochs if self._data.ndim == 4 else Evoked
+        unknown_class = Epochs if "epoch" in self._dims else Evoked
         inst_types = dict(Raw=Raw, Epochs=Epochs, Evoked=Evoked, Unknown=unknown_class)
         self._inst_type = inst_types[defaults["inst_type_str"]]
         # sanity check data/freqs/times/info agreement
@@ -1499,7 +1499,9 @@ class BaseTFR(ContainsMixin, UpdateChannelsMixin, SizeMixin, ExtendedTimeMixin):
         """Check data/freqs/times/info agreement during __setstate__."""
         msg = "{} axis of data ({}) doesn't match {} attribute ({})"
         n_chan_info = len(self.info["chs"])
-        n_chan, n_freq, n_time = self._data.shape[self._dims.index("channel") :]
+        n_chan = self._data.shape[self._dims.index("channel")]
+        n_freq = self._data.shape[self._dims.index("freq")]
+        n_time = self._data.shape[self._dims.index("time")]
         if n_chan_info != n_chan:
             msg = msg.format("Channel", n_chan, "info", n_chan_info)
         elif n_freq != len(self.freqs):
@@ -3372,13 +3374,14 @@ class EpochsTFR(BaseTFR, GetEpochsMixin):
                 "EpochsTFR.average() got a method that resulted in data of shape "
                 f"{data.shape}, but it should be {expected_shape}."
             )
+        state = self.__getstate__()
         # restore singleton freqs axis (not necessary for epochs/times: class changes)
         if dim == "freqs":
             data = np.expand_dims(data, axis=axis)
-        state = self.__getstate__()
+        else:
+            state["dims"] = (*state["dims"][:axis], *state["dims"][axis + 1 :])
         state["data"] = data
         state["info"] = deepcopy(self.info)
-        state["dims"] = (*state["dims"][:axis], *state["dims"][axis + 1 :])
         state["freqs"] = freqs
         state["times"] = times
         if dim == "epochs":
@@ -4211,7 +4214,7 @@ def read_tfrs(fname, condition=None, *, verbose=None):
     hdf5_dict = read_hdf5(fname, title="mnepython", slash="replace")
     # single TFR from TFR.save()
     if "inst_type_str" in hdf5_dict:
-        if hdf5_dict["data"].ndim == 4:
+        if "epoch" in hdf5_dict["dims"]:
             Klass = EpochsTFR
         elif "nave" in hdf5_dict:
             Klass = AverageTFR


Test patch:
diff --git a/mne/time_frequency/tests/test_psd.py b/mne/time_frequency/tests/test_psd.py
index e02e56138..363a4207c 100644
--- 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
index 752e1d000..a6ea0be97 100644
--- a/mne/time_frequency/tests/test_spectrum.py
+++ b/mne/time_frequency/tests/test_spectrum.py
@@ -24,6 +24,9 @@ def test_compute_psd_errors(raw):
         raw.compute_psd(foo=None, bar=None)
     with pytest.raises(ValueError, match="Complex output is not supported in "):
         raw.compute_psd(output="complex")
+    raw.set_annotations(Annotations(onset=0.01, duration=0.01, description="bad_foo"))
+    with pytest.raises(NotImplementedError, match='Cannot use method="multitaper"'):
+        raw.compute_psd(method="multitaper", reject_by_annotation=True)
 
 
 @pytest.mark.parametrize("method", ("welch", "multitaper"))
@@ -163,12 +166,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
diff --git a/mne/time_frequency/tests/test_tfr.py b/mne/time_frequency/tests/test_tfr.py
index 33cfdb96b..cedc13a47 100644
--- a/mne/time_frequency/tests/test_tfr.py
+++ b/mne/time_frequency/tests/test_tfr.py
@@ -1402,7 +1402,7 @@ def test_to_data_frame_time_format(time_format):
 @parametrize_morlet_multitaper
 @parametrize_power_phase_complex
 @pytest.mark.parametrize("picks", ("mag", mag_names, [2, 5, 8]))  # all 3 equivalent
-def test_raw_compute_tfr(raw, method, output, picks):
+def test_raw_compute_tfr(raw, method, output, picks, tmp_path):
     """Test Raw.compute_tfr() and picks handling."""
     full_tfr = raw.compute_tfr(method, output=output, freqs=freqs_linspace)
     pick_tfr = raw.compute_tfr(method, output=output, freqs=freqs_linspace, picks=picks)
@@ -1411,6 +1411,12 @@ def test_raw_compute_tfr(raw, method, output, picks):
     want = full_tfr.get_data(picks=mag_names)
     got = pick_tfr.get_data()
     assert_array_equal(want, got)
+    # make sure save/load works for phase/complex data
+    if output in ("phase", "complex"):
+        pytest.importorskip("h5io")
+        fname = tmp_path / "temp_tfr.hdf5"
+        full_tfr.save(fname, overwrite=True)
+        assert read_tfrs(fname) == full_tfr
 
 
 @parametrize_morlet_multitaper