Epochs.compute_tfr() fails for complex and phase outputs in the multitaper method

[tsbinns]

Description of the problem

When calling Epochs.compute_tfr() with method="multitaper" and output="complex"/"phase", there is an AssertionError that the actual shape of the results does not match the expected shape (self._shape):

The expected shape is incorrect due to the tapers dimension of the multitaper complex/phase outputs not being properly indexed in the epoched data.

There are no errors when TFR is computed from epoched data without a tapers dimension (e.g., multitaper mode with power output; Morlet mode with any output), or from raw data with a tapers dimension (multitaper mode with complex or phase outputs).

self._shape is generated here:

mne-python/mne/time_frequency/tfr.py

Lines 1527 to 1537 in 013d6c9

	# this is *expected* shape, it gets asserted later in _check_values()
	# (and then deleted afterwards)
	expected_shape = [
	len(self.ch_names),
	len(self.freqs),
	len(self._raw_times[self._decim]), # don't use self.times, not set yet
	]
	# deal with the "taper" dimension
	if self._needs_taper_dim:
	expected_shape.insert(1, self._data.shape[1])
	self._shape = tuple(expected_shape)

TFR results with tapers have shape ([epochs, ]channels, tapers, freqs, times). When tapers are present, they are added to expected_shape in L1536 with self._data.shape[1], however this only corresponds to the tapers dimension in non-epoched data. If the data is epoched, it corresponds to the number of channels, so self._shape becomes (epochs, channels, channels, freqs, times). This causes the mismatch between the actual and expected shapes of the results.

This could be fixed with something like:

# deal with the "taper" dimension
if self._needs_taper_dim:
    tapers_dim = 1 if _get_instance_type_string(self) != "Epochs" else 2
    expected_shape.insert(1, self._data.shape[tapers_dim])

Inserting the number of tapers into expected_shape in position 1 is fine for both epoched and non-epoched data, since the number of epochs gets added later:

mne-python/mne/time_frequency/tfr.py

Line 3155 in 013d6c9

self._shape = (len(self.inst),) + self._shape

Happy to open a PR for this.

Steps to reproduce

import traceback

import numpy as np
from mne import EpochsArray, create_info

# Generate data
rng = np.random.default_rng(44)
n_epochs, n_chans, n_times = (5, 2, 100)
sfreq = 50
data = rng.random((n_epochs, n_chans, n_times))
info = create_info(ch_names=n_chans, sfreq=sfreq, ch_types="eeg")
epochs = EpochsArray(data=data, info=info)
freqs = np.arange(10, sfreq // 2)

# Compute TFR
for output in ["power", "complex", "phase"]:
    try:
        tfr = epochs.compute_tfr(method="multitaper", freqs=freqs, output=output)
        print(f"\nOutput {output} passed.")
    except AssertionError:
        print(f"\nOutput {output} failed:")
        traceback.print_exc()

Link to data

No response

Expected results

No AssertionError should be raised for any of the outputs.

Actual results

An AssertionError is raised for "complex" and "phase" outputs, but not "power":

Additional information

Platform Windows-11-10.0.22631-SP0
Python 3.12.5 | packaged by conda-forge | (main, Aug 8 2024, 18:24:51) [MSC v.1940 64 bit (AMD64)]
Executable c:\Users\tsbin\anaconda3\envs\mne\python.exe
CPU Intel64 Family 6 Model 154 Stepping 3, GenuineIntel (16 cores)
Memory 31.7 GB

Core
├☑ mne 1.9.0.dev20+g013d6c9f6 (devel, latest release is 1.8.0)
├☑ numpy 1.26.4 (OpenBLAS 0.3.23.dev with 16 threads)
├☑ scipy 1.14.1
└☑ matplotlib 3.9.2 (backend=qtagg)