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:
Hello! 👋 Thanks for opening your first issue here! ❤️ We will try to get back to you soon. 🚴


Resolving patch:
diff --git a/mne/time_frequency/tfr.py b/mne/time_frequency/tfr.py
--- a/mne/time_frequency/tfr.py
+++ b/mne/time_frequency/tfr.py
@@ -1440,7 +1440,7 @@ def __setstate__(self, state):
         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 @@ def _check_state(self):
         """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 @@ def average(self, method="mean", *, dim="epochs", copy=False):
                 "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_tfr.py b/mne/time_frequency/tests/test_tfr.py
--- 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