Repository: mne-tools/mne-python

Issue:
tfr_array_multitaper with output="complex" is misleading
#### Describe the bug
the output of tfr_array_multitaper when output='complex' is misleading. Not sure if one should change something in the design, or just update the documentation.


#### Steps to reproduce
```
import mne.time_frequency as tfr
import numpy as np


signal = np.zeros(1000)
signal[500:] = 1

# morlet wavelets
tfr_out1 = tfr.tfr_array_morlet([[signal]], sfreq=100, freqs=[20], output='complex')
tfr_out2 = tfr.tfr_array_morlet([[signal]], sfreq=100, freqs=[20], output='power')
print(np.allclose(np.abs(tfr_out1)**2, tfr_out2))
print(tfr_out1.shape, tfr_out2.shape)

# using only 1 taper
tfr_out1 = tfr.tfr_array_multitaper([[signal]], sfreq=100, freqs=[20], time_bandwidth=2, output='complex')
tfr_out2 = tfr.tfr_array_multitaper([[signal]], sfreq=100, freqs=[20], time_bandwidth=2, output='power')
print(np.allclose(np.abs(tfr_out1)**2, tfr_out2))
print(tfr_out1.shape, tfr_out2.shape)

# using only 2 tapers
tfr_out1 = tfr.tfr_array_multitaper([[signal]], sfreq=100, freqs=[20], time_bandwidth=3, output='complex')
tfr_out2 = tfr.tfr_array_multitaper([[signal]], sfreq=100, freqs=[20], time_bandwidth=3, output='power')
print(np.allclose(np.abs(tfr_out1)**2, tfr_out2))
print(tfr_out1.shape, tfr_out2.shape)
```

#### Expected vs Actual results
Everything is as expected, apart from the last output: when using multitaper with more than 1 taper, one should get multiple complex time series outputs, one for each taper, so that one can reconstruct the power from the complex time series. With the current design, it is not possible, and it is not clear what is the meaning of the output returned by tfr.tfr_array_multitaper([[signal]], sfreq=100, freqs=[20], time_bandwidth=3, output='complex'). Is it the convolution with one of the tapers? Is it a sum of the complex time series from different tapers?


#### Actual results
One gets:

Morlet:
True
(1, 1, 1, 1000) (1, 1, 1, 1000)

multitaper with single taper:
True 
(1, 1, 1, 1000) (1, 1, 1, 1000)

multitaper with 2 tapers:
False 
(1, 1, 1, 1000) (1, 1, 1, 1000) (we need more information than that to reconstruct the "power" from "complex")


#### Additional information
Platform:      macOS-10.15.7-x86_64-i386-64bit
Python:        3.8.5 (default, Aug  5 2020, 03:39:04)  [Clang 10.0.0 ]
Executable:    /Users/.../opt/anaconda3/envs/mne/bin/python
CPU:           i386: 12 cores
Memory:        16.0 GB
mne:           0.22.0
numpy:         1.19.4 {blas=NO_ATLAS_INFO, lapack=lapack}
scipy:         1.5.3
matplotlib:    3.3.1 {backend=MacOSX}
sklearn:       0.24.0
numba:         Not found
nibabel:       3.2.1
nilearn:       0.7.0
dipy:          1.3.0
cupy:          Not found
pandas:        1.1.5
mayavi:        4.7.2
pyvista:       0.27.4 {pyvistaqt=0.2.0, OpenGL 4.1 ATI-3.10.19 via AMD Radeon Pro 560X OpenGL Engine}
vtk:           9.0.1
PyQt5:         5.12.3




Discussion:
Let's fix this with https://github.com/mne-tools/mne-python/issues/8724 in
one PR

>


Resolving patch:
diff --git a/mne/time_frequency/multitaper.py b/mne/time_frequency/multitaper.py
--- a/mne/time_frequency/multitaper.py
+++ b/mne/time_frequency/multitaper.py
@@ -494,9 +494,9 @@ def tfr_array_multitaper(epoch_data, sfreq, freqs, n_cycles=7.0,
             is done after the convolutions.
     output : str, default 'complex'
 
-        * 'complex' : single trial complex.
+        * 'complex' : single trial per taper complex values.
         * 'power' : single trial power.
-        * 'phase' : single trial phase.
+        * 'phase' : single trial per taper phase.
         * 'avg_power' : average of single trial power.
         * 'itc' : inter-trial coherence.
         * 'avg_power_itc' : average of single trial power and inter-trial
@@ -509,11 +509,13 @@ def tfr_array_multitaper(epoch_data, sfreq, freqs, n_cycles=7.0,
     Returns
     -------
     out : array
-        Time frequency transform of epoch_data. If output is in ['complex',
-        'phase', 'power'], then shape of out is (n_epochs, n_chans, n_freqs,
-        n_times), else it is (n_chans, n_freqs, n_times). If output is
-        'avg_power_itc', the real values code for 'avg_power' and the
-        imaginary values code for the 'itc': out = avg_power + i * itc.
+        Time frequency transform of epoch_data. If ``output in ['complex',
+        'phase']``, then the shape of ``out`` is ``(n_epochs, n_chans,
+        n_tapers, n_freqs, n_times)``; if output is 'power', the shape of
+        ``out`` is ``(n_epochs, n_chans, n_freqs, n_times)``, else it is
+        ``(n_chans, n_freqs, n_times)``. If output is 'avg_power_itc', the real
+        values in ``out`` contain the average power and the imaginary values
+        contain the ITC: ``out = avg_power + i * itc``.
 
     See Also
     --------
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
@@ -331,10 +331,13 @@ def _compute_tfr(epoch_data, freqs, sfreq=1.0, method='morlet',
     -------
     out : array
         Time frequency transform of epoch_data. If output is in ['complex',
-        'phase', 'power'], then shape of out is (n_epochs, n_chans, n_freqs,
-        n_times), else it is (n_chans, n_freqs, n_times). If output is
-        'avg_power_itc', the real values code for 'avg_power' and the
-        imaginary values code for the 'itc': out = avg_power + i * itc
+        'phase', 'power'], then shape of ``out`` is ``(n_epochs, n_chans,
+        n_freqs, n_times)``, else it is ``(n_chans, n_freqs, n_times)``.
+        However, using multitaper method and output ``'complex'`` or
+        ``'phase'`` results in shape of ``out`` being ``(n_epochs, n_chans,
+        n_tapers, n_freqs, n_times)``. If output is ``'avg_power_itc'``, the
+        real values in the ``output`` contain average power' and the imaginary
+        values contain the ITC: ``out = avg_power + i * itc``.
     """
     # Check data
     epoch_data = np.asarray(epoch_data)
@@ -370,6 +373,7 @@ def _compute_tfr(epoch_data, freqs, sfreq=1.0, method='morlet',
 
     # Initialize output
     n_freqs = len(freqs)
+    n_tapers = len(Ws)
     n_epochs, n_chans, n_times = epoch_data[:, :, decim].shape
     if output in ('power', 'phase', 'avg_power', 'itc'):
         dtype = np.float64
@@ -380,6 +384,8 @@ def _compute_tfr(epoch_data, freqs, sfreq=1.0, method='morlet',
 
     if ('avg_' in output) or ('itc' in output):
         out = np.empty((n_chans, n_freqs, n_times), dtype)
+    elif output in ['complex', 'phase'] and method == 'multitaper':
+        out = np.empty((n_chans, n_tapers, n_epochs, n_freqs, n_times), dtype)
     else:
         out = np.empty((n_chans, n_epochs, n_freqs, n_times), dtype)
 
@@ -390,7 +396,7 @@ def _compute_tfr(epoch_data, freqs, sfreq=1.0, method='morlet',
 
     # Parallelization is applied across channels.
     tfrs = parallel(
-        my_cwt(channel, Ws, output, use_fft, 'same', decim)
+        my_cwt(channel, Ws, output, use_fft, 'same', decim, method)
         for channel in epoch_data.transpose(1, 0, 2))
 
     # FIXME: to avoid overheads we should use np.array_split()
@@ -399,7 +405,10 @@ def _compute_tfr(epoch_data, freqs, sfreq=1.0, method='morlet',
 
     if ('avg_' not in output) and ('itc' not in output):
         # This is to enforce that the first dimension is for epochs
-        out = out.transpose(1, 0, 2, 3)
+        if output in ['complex', 'phase'] and method == 'multitaper':
+            out = out.transpose(2, 0, 1, 3, 4)
+        else:
+            out = out.transpose(1, 0, 2, 3)
     return out
 
 
@@ -428,11 +437,6 @@ def _check_tfr_param(freqs, sfreq, method, zero_mean, n_cycles,
                          % type(zero_mean))
     freqs = np.asarray(freqs)
 
-    if (method == 'multitaper') and (output == 'phase'):
-        raise NotImplementedError(
-            'This function is not optimized to compute the phase using the '
-            'multitaper method. Use np.angle of the complex output instead.')
-
     # Check n_cycles
     if isinstance(n_cycles, (int, float)):
         n_cycles = float(n_cycles)
@@ -472,7 +476,8 @@ def _check_tfr_param(freqs, sfreq, method, zero_mean, n_cycles,
     return freqs, sfreq, zero_mean, n_cycles, time_bandwidth, decim
 
 
-def _time_frequency_loop(X, Ws, output, use_fft, mode, decim):
+def _time_frequency_loop(X, Ws, output, use_fft, mode, decim,
+                         method=None):
     """Aux. function to _compute_tfr.
 
     Loops time-frequency transform across wavelets and epochs.
@@ -499,6 +504,9 @@ def _time_frequency_loop(X, Ws, output, use_fft, mode, decim):
         See numpy.convolve.
     decim : slice
         The decimation slice: e.g. power[:, decim]
+    method : str | None
+        Used only for multitapering to create tapers dimension in the output
+        if ``output in ['complex', 'phase']``.
     """
     # Set output type
     dtype = np.float64
@@ -507,15 +515,19 @@ def _time_frequency_loop(X, Ws, output, use_fft, mode, decim):
 
     # Init outputs
     decim = _check_decim(decim)
+    n_tapers = len(Ws)
     n_epochs, n_times = X[:, decim].shape
     n_freqs = len(Ws[0])
     if ('avg_' in output) or ('itc' in output):
         tfrs = np.zeros((n_freqs, n_times), dtype=dtype)
+    elif output in ['complex', 'phase'] and method == 'multitaper':
+        tfrs = np.zeros((n_tapers, n_epochs, n_freqs, n_times),
+                        dtype=dtype)
     else:
         tfrs = np.zeros((n_epochs, n_freqs, n_times), dtype=dtype)
 
     # Loops across tapers.
-    for W in Ws:
+    for taper_idx, W in enumerate(Ws):
         # No need to check here, it's done earlier (outside parallel part)
         nfft = _get_nfft(W, X, use_fft, check=False)
         coefs = _cwt_gen(
@@ -543,6 +555,8 @@ def _time_frequency_loop(X, Ws, output, use_fft, mode, decim):
             # Stack or add
             if ('avg_' in output) or ('itc' in output):
                 tfrs += tfr
+            elif output in ['complex', 'phase'] and method == 'multitaper':
+                tfrs[taper_idx, epoch_idx] += tfr
             else:
                 tfrs[epoch_idx] += tfr
 
@@ -557,7 +571,8 @@ def _time_frequency_loop(X, Ws, output, use_fft, mode, decim):
         tfrs /= n_epochs
 
     # Normalization by number of taper
-    tfrs /= len(Ws)
+    if n_tapers > 1 and output not in ['complex', 'phase']:
+        tfrs /= n_tapers
     return tfrs
 
 


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
@@ -131,6 +131,21 @@ def test_time_frequency():
     # computed within the method.
     assert_allclose(epochs_amplitude_2.data**2, epochs_power_picks.data)
 
+    # test that averaging power across tapers when multitaper with
+    # output='complex' gives the same as output='power'
+    epoch_data = epochs.get_data()
+    multitaper_power = tfr_array_multitaper(
+        epoch_data, epochs.info['sfreq'], freqs, n_cycles,
+        output="power")
+    multitaper_complex = tfr_array_multitaper(
+        epoch_data, epochs.info['sfreq'], freqs, n_cycles,
+        output="complex")
+
+    taper_dim = 2
+    power_from_complex = (multitaper_complex * multitaper_complex.conj()
+                          ).real.mean(axis=taper_dim)
+    assert_allclose(power_from_complex, multitaper_power)
+
     print(itc)  # test repr
     print(itc.ch_names)  # test property
     itc += power  # test add
@@ -721,17 +736,16 @@ def test_compute_tfr():
         (tfr_array_multitaper, tfr_array_morlet), (False, True), (False, True),
         ('complex', 'power', 'phase',
          'avg_power_itc', 'avg_power', 'itc')):
-        # Check exception
-        if (func == tfr_array_multitaper) and (output == 'phase'):
-            pytest.raises(NotImplementedError, func, data, sfreq=sfreq,
-                          freqs=freqs, output=output)
-            continue
 
         # Check runs
         out = func(data, sfreq=sfreq, freqs=freqs, use_fft=use_fft,
                    zero_mean=zero_mean, n_cycles=2., output=output)
         # Check shapes
-        shape = np.r_[data.shape[:2], len(freqs), data.shape[2]]
+        if func == tfr_array_multitaper and output in ['complex', 'phase']:
+            n_tapers = 3
+            shape = np.r_[data.shape[:2], n_tapers, len(freqs), data.shape[2]]
+        else:
+            shape = np.r_[data.shape[:2], len(freqs), data.shape[2]]
         if ('avg' in output) or ('itc' in output):
             assert_array_equal(shape[1:], out.shape)
         else:
@@ -762,9 +776,6 @@ def test_compute_tfr():
     # No time_bandwidth param in morlet
     pytest.raises(ValueError, _compute_tfr, data, freqs, sfreq,
                   method='morlet', time_bandwidth=1)
-    # No phase in multitaper XXX Check ?
-    pytest.raises(NotImplementedError, _compute_tfr, data, freqs, sfreq,
-                  method='multitaper', output='phase')
 
     # Inter-trial coherence tests
     out = _compute_tfr(data, freqs, sfreq, output='itc', n_cycles=2.)
@@ -780,10 +791,11 @@ def test_compute_tfr():
         _decim = slice(None, None, decim) if isinstance(decim, int) else decim
         n_time = len(np.arange(data.shape[2])[_decim])
         shape = np.r_[data.shape[:2], len(freqs), n_time]
+
         for method in ('multitaper', 'morlet'):
             # Single trials
             out = _compute_tfr(data, freqs, sfreq, method=method, decim=decim,
-                               n_cycles=2.)
+                               output='power', n_cycles=2.)
             assert_array_equal(shape, out.shape)
             # Averages
             out = _compute_tfr(data, freqs, sfreq, method=method, decim=decim,
@@ -798,14 +810,17 @@ def test_compute_tfr_correct(method, decim):
     sfreq = 1000.
     t = np.arange(1000) / sfreq
     f = 50.
-    data = np.sin(2 * np.pi * 50. * t)
+    data = np.sin(2 * np.pi * f * t)
     data *= np.hanning(data.size)
     data = data[np.newaxis, np.newaxis]
-    freqs = np.arange(10, 111, 10)
+    freqs = np.arange(10, 111, 4)
     assert f in freqs
+
+    # previous n_cycles=2 gives weird results for multitaper
+    n_cycles = freqs * 0.25
     tfr = _compute_tfr(data, freqs, sfreq, method=method, decim=decim,
-                       n_cycles=2)[0, 0]
-    assert freqs[np.argmax(np.abs(tfr).mean(-1))] == f
+                       n_cycles=n_cycles, output='power')[0, 0]
+    assert freqs[np.argmax(tfr.mean(-1))] == f
 
 
 def test_averaging_epochsTFR():