Support complex output in EpochsSpectrum Welch (#11556)

Author: larsoner

doc/changes/latest.inc

@@ -32,6 +32,7 @@ Enhancements
 - Add :meth:`mne.Info.save` to save an :class:`mne.Info` object to a fif file (:gh:`11401` by `Alex Rockhill`_)
 - Improved error message when downloads are corrupted for :func:`mne.datasets.sample.data_path` and related functions (:gh:`11407` by `Eric Larson`_)
 - Add support for ``skip_by_annotation`` in :func:`mne.io.Raw.notch_filter` (:gh:`11388` by `Mainak Jas`_)
+- Add support for ``output='complex'`` to :func:`mne.time_frequency.psd_array_welch` and when using ``method='welch'`` with :meth:`mne.Epochs.compute_psd` (:gh:`11556` by `Eric Larson`_)
 - Slightly adjusted the window title for :func:`mne.Epochs.plot` (:gh:`11419` by `Richard Höchenberger`_ and `Daniel McCloy`_)
 - Add :func:`mne.count_events` to count unique event types in a given event array (:gh:`11430` by `Clemens Brunner`_)
 - Add a video to :ref:`tut-freesurfer-mne` of a brain inflating from the pial surface to aid in understanding the inflated brain (:gh:`11440` by `Alex Rockhill`_)

mne/conftest.py

@@ -131,6 +131,7 @@ def pytest_configure(config):
     ignore:Jupyter is migrating its paths.*:DeprecationWarning
     ignore:Widget\..* is deprecated\.:DeprecationWarning
     ignore:.*is deprecated in pyzmq.*:DeprecationWarning
+    ignore:The `ipykernel.comm.Comm` class has been deprecated.*:DeprecationWarning
     # PySide6
     ignore:Enum value .* is marked as deprecated:DeprecationWarning
     ignore:Function.*is marked as deprecated, please check the documentation.*:DeprecationWarning

mne/time_frequency/psd.py

@@ -56,7 +56,7 @@ def _decomp_aggregate_mask(epoch, func, average, freq_sl):
     return spect
 
 
-def _spect_func(epoch, func, freq_sl, average):
+def _spect_func(epoch, func, freq_sl, average, *, output='power'):
     """Aux function."""
     # Decide if we should split this to save memory or not, since doing
     # multiple calls will incur some performance overhead. Eventually we might
@@ -91,7 +91,7 @@ def _check_nfft(n, n_fft, n_per_seg, n_overlap):
 @verbose
 def psd_array_welch(x, sfreq, fmin=0, fmax=np.inf, n_fft=256, n_overlap=0,
                     n_per_seg=None, n_jobs=None, average='mean',
-                    window='hamming', *, verbose=None):
+                    window='hamming', *, output='power', verbose=None):
     """Compute power spectral density (PSD) using Welch's method.
 
     Welch's method is described in :footcite:t:`Welch1967`.
@@ -122,6 +122,15 @@ def psd_array_welch(x, sfreq, fmin=0, fmax=np.inf, n_fft=256, n_overlap=0,
     %(window_psd)s
 
         .. versionadded:: 0.22.0
+    output : str
+        The format of the returned ``psds`` array, ``'complex'`` or
+        ``'power'``:
+
+        * ``'power'`` : the power spectral density is returned.
+        * ``'complex'`` : the complex fourier coefficients are returned per
+          window.
+
+        .. versionadded:: 1.4.0
     %(verbose)s
 
     Returns
@@ -145,6 +154,8 @@ def psd_array_welch(x, sfreq, fmin=0, fmax=np.inf, n_fft=256, n_overlap=0,
     .. footbibliography::
     """
     _check_option('average', average, (None, False, 'mean', 'median'))
+    _check_option('output', output, ('power', 'complex'))
+    mode = 'complex' if output == 'complex' else 'psd'
     n_fft = _ensure_int(n_fft, "n_fft")
     n_overlap = _ensure_int(n_overlap, "n_overlap")
     if n_per_seg is not None:
@@ -178,10 +189,10 @@ def psd_array_welch(x, sfreq, fmin=0, fmax=np.inf, n_fft=256, n_overlap=0,
     from scipy.signal import spectrogram
     parallel, my_spect_func, n_jobs = parallel_func(_spect_func, n_jobs=n_jobs)
     func = partial(spectrogram, noverlap=n_overlap, nperseg=n_per_seg,
-                   nfft=n_fft, fs=sfreq, window=window)
+                   nfft=n_fft, fs=sfreq, window=window, mode=mode)
     x_splits = [arr for arr in np.array_split(x, n_jobs) if arr.size != 0]
     f_spect = parallel(my_spect_func(d, func=func, freq_sl=freq_sl,
-                                     average=average)
+                                     average=average, output=output)
                        for d in x_splits)
     psds = np.concatenate(f_spect, axis=0)
     shape = dshape + (len(freqs),)

mne/time_frequency/spectrum.py

@@ -252,7 +252,7 @@ def __init__(self, inst, method, fmin, fmax, tmin, tmax, picks,
         self._dims = ('channel', 'freq',)
         if method_kw.get('average', '') in (None, False):
             self._dims += ('segment',)
-        if method_kw.get('output', '') == 'complex':
+        if self._returns_complex_tapers(**method_kw):
             self._dims = self._dims[:-1] + ('taper',) + self._dims[-1:]
         # record data type (for repr and html_repr)
         self._data_type = ('Fourier Coefficients' if 'taper' in self._dims
@@ -316,7 +316,8 @@ def _repr_html_(self, caption=None):
 
     def _check_values(self):
         """Check PSD results for correct shape and bad values."""
-        assert len(self._dims) == self._data.ndim
+        assert len(self._dims) == self._data.ndim, \
+            (self._dims, self._data.ndim)
         assert self._data.shape == self._shape
         # negative values OK if the spectrum is really fourier coefficients
         if 'taper' in self._dims:
@@ -333,13 +334,19 @@ def _check_values(self):
             warn(f'Zero value in spectrum for channel{s} {", ".join(chs)}',
                  UserWarning)
 
+    def _returns_complex_tapers(self, **method_kw):
+        return (
+            method_kw.get('output', '') == 'complex' and
+            self.method == 'multitaper'
+        )
+
     def _compute_spectra(self, data, fmin, fmax, n_jobs, method_kw, verbose):
         # make the spectra
         result = self._psd_func(
             data, self.sfreq, fmin=fmin, fmax=fmax, n_jobs=n_jobs,
             verbose=verbose)
         # assign ._data (handling unaggregated multitaper output)
-        if method_kw.get('output', '') == 'complex':
+        if self._returns_complex_tapers(**method_kw):
             fourier_coefs, freqs, weights = result
             self._data = fourier_coefs
             self._mt_weights = weights
@@ -357,7 +364,7 @@ def _compute_spectra(self, data, fmin, fmax, n_jobs, method_kw, verbose):
                                                          method_kw)
             self._shape += (n_welch_segments,)
         # insert n_tapers
-        if method_kw.get('output', '') == 'complex':
+        if self._returns_complex_tapers(**method_kw):
             self._shape = (
                 self._shape[:-1] + (self._mt_weights.size,) + self._shape[-1:])
         # we don't need these anymore, and they make save/load harder

mne/time_frequency/tests/test_spectrum.py

@@ -1,9 +1,12 @@
+from contextlib import nullcontext
 from functools import partial
 
 import numpy as np
 import pytest
-from numpy.testing import assert_array_equal
+from numpy.testing import assert_array_equal, assert_allclose
 
+from mne import create_info, make_fixed_length_epochs
+from mne.io import RawArray
 from mne import Annotations
 from mne.time_frequency import read_spectrum
 from mne.time_frequency.multitaper import _psd_from_mt
@@ -137,8 +140,12 @@ def _agg_helper(df, weights, group_cols):
 
 @requires_pandas
 @pytest.mark.parametrize('long_format', (False, True))
-@pytest.mark.parametrize('method', ('welch', 'multitaper'))
-def test_unaggregated_spectrum_to_data_frame(raw, long_format, method):
+@pytest.mark.parametrize('method, output', [
+    ('welch', 'complex'),
+    ('welch', 'power'),
+    ('multitaper', 'complex'),
+])
+def test_unaggregated_spectrum_to_data_frame(raw, long_format, method, output):
     """Test converting complex multitaper spectra to data frame."""
     from pandas.testing import assert_frame_equal
 
@@ -149,8 +156,10 @@ def test_unaggregated_spectrum_to_data_frame(raw, long_format, method):
                   .to_data_frame(long_format=long_format))
     # unaggregated welch or complex multitaper →
     #   aggregate w/ pandas (to make sure we did reshaping right)
-    kwargs = {'average': False} if method == 'welch' else {'output': 'complex'}
-    spectrum = raw.compute_psd(method=method, **kwargs)
+    kwargs = dict()
+    if method == 'welch':
+        kwargs.update(average=False, verbose='error')
+    spectrum = raw.compute_psd(method=method, output=output, **kwargs)
     df = spectrum.to_data_frame(long_format=long_format)
     grouping_cols = ['freq']
     drop_cols = ['segment'] if method == 'welch' else ['taper']
@@ -169,7 +178,12 @@ def test_unaggregated_spectrum_to_data_frame(raw, long_format, method):
     gb = df.drop(columns=drop_cols).groupby(
         grouping_cols, as_index=False, observed=False)
     if method == 'welch':
-        agg_df = gb.aggregate(np.nanmean)
+        if output == 'complex':
+            def _fun(x):
+                return np.nanmean(np.abs(x))
+        else:
+            _fun = np.nanmean
+        agg_df = gb.aggregate(_fun)
     else:
         agg_df = gb.apply(_agg_helper, spectrum._mt_weights, grouping_cols)
     # even with check_categorical=False, we know that the *data* matches;
@@ -235,3 +249,52 @@ def test_spectrum_proj(inst, request):
     with has_proj.info._unlock():
         has_proj.info['projs'] = no_proj.info['projs']
     assert has_proj == no_proj
+
+
+@pytest.mark.parametrize('method, average', [
+    ('welch', False),
+    ('welch', 'mean'),
+    ('multitaper', False),
+])
+def test_spectrum_complex(method, average):
+    """Test output='complex' support."""
+    sfreq = 100
+    n = 10 * sfreq
+    freq = 3.
+    phase = np.pi / 4  # should be recoverable
+    data = np.cos(2 * np.pi * freq * np.arange(n) / sfreq + phase)[np.newaxis]
+    raw = RawArray(data, create_info(1, sfreq, 'eeg'))
+    epochs = make_fixed_length_epochs(raw, duration=2., preload=True)
+    assert len(epochs) == 5
+    assert len(epochs.times) == 2 * sfreq
+    kwargs = dict(output='complex', method=method)
+    if method == 'welch':
+        kwargs['n_fft'] = sfreq
+        ctx = pytest.warns(UserWarning, match='Zero value')
+        want_dims = ('epoch', 'channel', 'freq')
+        want_shape = (5, 1, sfreq // 2 + 1)
+        if not average:
+            want_dims = want_dims + ('segment',)
+            want_shape = want_shape + (2,)
+            kwargs['average'] = average
+    else:
+        assert method == 'multitaper'
+        assert not average
+        ctx = nullcontext()
+        want_dims = ('epoch', 'channel', 'taper', 'freq')
+        want_shape = (5, 1, 7, sfreq + 1)
+    with ctx:
+        spectrum = epochs.compute_psd(**kwargs)
+    idx = np.argmin(np.abs(spectrum.freqs - freq))
+    assert spectrum.freqs[idx] == freq
+    assert spectrum._dims == want_dims
+    assert spectrum.shape == want_shape
+    data = spectrum.get_data()
+    assert data.dtype == np.complex128
+    coef = spectrum.get_data(fmin=freq, fmax=freq).mean(0)
+    if method == 'multitaper':
+        coef = coef[..., 0, :]  # first taper
+    elif not average:
+        coef = coef.mean(-1)  # over segments
+    coef = coef.item()
+    assert_allclose(np.angle(coef), phase, rtol=1e-4)