Merge remote-tracking branch 'upstream/main'

Author: tsbinns

changelog.md

@@ -4,6 +4,7 @@
 
 ##### Enhancements
 - Added `norm` parameter to `WaveShape.compute()` to control normalisation of waveshape results.
+- Added `output` parameter to `compute_tfr()` to allow complex coefficients to be returned.
 
 <br>
 

docs/source/conf.py

@@ -136,6 +136,7 @@
     "components",
     "frequencies",
     "frequency_bands",
+    "tapers",
     "x",
     "n_vertices",
     "n_faces",

src/pybispectra/utils/utils.py

@@ -8,7 +8,7 @@
 import pooch
 import numpy as np
 import scipy as sp
-from mne import time_frequency
+from mne import time_frequency, __version__ as mne_version
 
 from pybispectra import __version__ as version
 from pybispectra.utils._defaults import _precision
@@ -164,15 +164,16 @@ def compute_tfr(
     zero_mean_wavelets: bool | None = None,
     use_fft: bool = True,
     multitaper_time_bandwidth: int | float = 4.0,
+    output: str = "power",
     n_jobs: int = 1,
     verbose: bool = True,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Compute the amplitude time-frequency representation (TFR) of data.
+) -> tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Compute the time-frequency representation (TFR) of data.
 
     Parameters
     ----------
     data : ~numpy.ndarray, shape of [epochs, channels, times]
-        Real-valued data to compute the amplitude TFR of.
+        Real-valued data to compute the TFR of.
 
     sampling_freq : int | float
         Sampling frequency (in Hz) of ``data``.
@@ -203,6 +204,14 @@ def compute_tfr(
         bandwidth (in Hz). Only used if ``tfr_mode = "multitaper"``. See
         :func:`mne.time_frequency.tfr_array_multitaper` for more information.
 
+    output : ``"power"`` | ``"complex"`` (default ``"power"``)
+        Type of TFR output to return.
+
+        .. note::
+            If ``output = "complex"`` and ``tfr_mode = "multitaper"``, returning weights
+            for each taper requires MNE version 1.10 or higher.
+        .. versionadded:: 1.3
+
     n_jobs : int (default ``1``)
         Number of jobs to run in parallel. If ``-1``, all available CPUs are used.
 
@@ -211,19 +220,24 @@ def compute_tfr(
 
     Returns
     -------
-    tfr : ~numpy.ndarray, shape of [epochs, channels, frequencies, times]
-        Amplitude/power of the TFR of ``data``.
+    tfr : ~numpy.ndarray, shape of [epochs, channels (, tapers), frequencies, times]
+        TFR power or complex coefficients of ``data``. The ``tapers`` dimension is only
+        present if ``output = "complex"`` and ``tfr_mode = "multitaper"``.
 
     freqs : ~numpy.ndarray of float, shape of [frequencies]
         Frequencies (in Hz) in ``tfr``.
 
+    weights : ~numpy.ndarray, shape of [tapers, frequencies]
+        Taper weights. Only returned if ``output = "complex"`` and ``tfr_mode =
+        "multitaper"``.
+
     Notes
     -----
     This function acts as a wrapper around the MNE TFR computation functions
     :func:`mne.time_frequency.tfr_array_morlet` and
-    :func:`mne.time_frequency.tfr_array_multitaper` with ``output = "power"``.
+    :func:`mne.time_frequency.tfr_array_multitaper`.
     """
-    tfr_func, n_jobs = _compute_tfr_input_checks(
+    tfr_func, return_weights, n_jobs = _compute_tfr_input_checks(
         data,
         sampling_freq,
         freqs,
@@ -232,6 +246,7 @@ def compute_tfr(
         zero_mean_wavelets,
         use_fft,
         multitaper_time_bandwidth,
+        output,
         n_jobs,
         verbose,
     )
@@ -242,23 +257,35 @@ def compute_tfr(
         "freqs": freqs,
         "n_cycles": n_cycles,
         "use_fft": use_fft,
-        "output": "power",
+        "output": output,
         "n_jobs": n_jobs,
         "verbose": verbose,
     }
     if zero_mean_wavelets is not None:
         tfr_func_kwargs["zero_mean"] = zero_mean_wavelets
     if tfr_mode == "multitaper":
         tfr_func_kwargs["time_bandwidth"] = multitaper_time_bandwidth
+        if output == "complex":
+            tfr_func_kwargs["return_weights"] = True
 
     if verbose:
         print("Computing TFR of the data...")
 
-    tfr = np.array(tfr_func(**tfr_func_kwargs), dtype=_precision.real)
+    out = tfr_func(**tfr_func_kwargs)
+    if return_weights:
+        tfr = out[0]
+        weights = out[1]
+    else:
+        tfr = out
+    tfr = np.asarray(
+        tfr, dtype=_precision.real if output == "power" else _precision.complex
+    )
 
     if verbose:
         print("    [TFR computation finished]\n")
 
+    if return_weights:
+        return tfr, freqs.astype(_precision.real), weights.astype(_precision.real)
     return tfr, freqs.astype(_precision.real)
 
 
@@ -271,16 +298,20 @@ def _compute_tfr_input_checks(
     zero_mean_wavelets: bool | None,
     use_fft: bool,
     multitaper_time_bandwidth: int | float,
+    output: str,
     n_jobs: int,
     verbose: bool,
-) -> tuple[Callable, int]:
+) -> tuple[Callable, bool, int]:
     """Check inputs for computing TFR.
 
     Returns
     -------
     tfr_func
         Function to use to compute TFR.
 
+    return_weights : bool
+        Whether or not taper weights will be returned.
+
     n_jobs
     """
     if not isinstance(data, np.ndarray):
@@ -334,6 +365,21 @@ def _compute_tfr_input_checks(
         if not isinstance(multitaper_time_bandwidth, _number_like):
             raise TypeError("`multitaper_time_bandwidth` must be an int or a float.")
 
+    outputs = ["power", "complex"]
+    if not isinstance(output, str):
+        raise TypeError("`output` must be a str.")
+    if output not in outputs:
+        raise ValueError(f"`output` must be one of {outputs}.")
+
+    return_weights = False
+    if tfr_mode == "multitaper" and output == "complex":  # pragma: no cover
+        if Version(mne_version) < Version("1.10"):
+            raise RuntimeError(
+                "If `tfr_mode='multitaper'` and `output='complex'`, MNE >= 1.10 is "
+                f"required to return taper weights (is {mne_version})."
+            )
+        return_weights = True
+
     if not isinstance(n_jobs, _int_like):
         raise TypeError("`n_jobs` must be an integer.")
     if n_jobs < 1 and n_jobs != -1:
@@ -346,7 +392,7 @@ def _compute_tfr_input_checks(
     if verbose and not np.isreal(data).all():
         warn("`data` is expected to be real-valued.", UserWarning)
 
-    return tfr_func, n_jobs
+    return tfr_func, return_weights, n_jobs
 
 
 def compute_rank(data: np.ndarray, sv_tol: int | float = 1e-5) -> int:

tests/test_util_funcs.py

@@ -1,11 +1,12 @@
 """Tests for toolbox utility functions."""
 
 import os
+from packaging.version import Version
 
 import numpy as np
 import pytest
 import scipy as sp
-from mne import Info
+from mne import Info, __version__ as mne_version
 
 from pybispectra.utils import (
     compute_fft,
@@ -106,9 +107,21 @@ def test_compute_fft(window: str) -> None:
 
 
 @pytest.mark.parametrize("tfr_mode", ["morlet", "multitaper"])
+@pytest.mark.parametrize("output", ["power", "complex"])
 @pytest.mark.parametrize("zero_mean_wavelets", [True, False, None])
-def test_compute_tfr(tfr_mode: str, zero_mean_wavelets: bool | None) -> None:
+def test_compute_tfr(
+    tfr_mode: str, output: str, zero_mean_wavelets: bool | None
+) -> None:
     """Test `compute_tfr`."""
+    if (
+        tfr_mode == "multitaper"
+        and output == "complex"
+        and Version(mne_version) < Version("1.10")
+    ):
+        pytest.skip(
+            "`output='complex'` with `tfr_mode='multitaper'` requires MNE >= 1.10."
+        )
+
     n_epochs = 5
     n_chans = 3
     n_times = 100
@@ -117,25 +130,41 @@ def test_compute_tfr(tfr_mode: str, zero_mean_wavelets: bool | None) -> None:
     freqs_in = np.arange(20, 50)
 
     # check it runs with correct inputs
-    tfr, freqs_out = compute_tfr(
+    out = compute_tfr(
         data=data,
         sampling_freq=sampling_freq,
         freqs=freqs_in,
         tfr_mode=tfr_mode,
         zero_mean_wavelets=zero_mean_wavelets,
+        output=output,
         n_jobs=1,
     )
+    if tfr_mode == "multitaper" and output == "complex":
+        tfr, freqs_out, weights = out
+    else:
+        tfr, freqs_out = out
+        weights = None
+
     assert isinstance(tfr, np.ndarray), "`tfr` should be a NumPy array."
-    assert tfr.ndim == 4, "`tfr` should have 4 dimensions."
-    assert tfr.shape[:3] == (n_epochs, n_chans, len(freqs_in)), (
-        "The first 3 dimensions of `tfr` should have shape [epochs x channels x "
-        "frequencies]."
-    )
+    if tfr_mode == "multitaper" and output == "complex":
+        assert tfr.shape == (n_epochs, n_chans, 3, len(freqs_in), n_times), (
+            "`tfr` should have shape [epochs x channels x tapers x frequencies x times]."
+        )  # unsure how to predict number of tapers, so use known value for this test
+    else:
+        assert tfr.shape == (n_epochs, n_chans, len(freqs_in), n_times), (
+            "`tfr` should have shape [epochs x channels x frequencies x times]."
+        )
     assert isinstance(freqs_out, np.ndarray), "`freqs_out` should be a NumPy array."
     assert np.all(freqs_in == freqs_out), (
         "`freqs_out` and `freqs_in` should be identical"
     )
 
+    if weights is not None:
+        assert isinstance(weights, np.ndarray), "`weights` should be a NumPy array."
+        assert weights.shape == tfr.shape[2:4], (
+            "`weights` should have shape [tapers x frequencies]."
+        )
+
     # check it catches incorrect inputs
     with pytest.raises(TypeError, match="`data` must be a NumPy array."):
         compute_tfr(
@@ -297,6 +326,25 @@ def test_compute_tfr(tfr_mode: str, zero_mean_wavelets: bool | None) -> None:
                 multitaper_time_bandwidth=[3],
             )
 
+    with pytest.raises(TypeError, match="`output` must be a str."):
+        compute_tfr(
+            data=data,
+            sampling_freq=sampling_freq,
+            freqs=freqs_in,
+            tfr_mode=tfr_mode,
+            zero_mean_wavelets=zero_mean_wavelets,
+            output=[output],
+        )
+    with pytest.raises(ValueError, match="`output` must be one of"):
+        compute_tfr(
+            data=data,
+            sampling_freq=sampling_freq,
+            freqs=freqs_in,
+            tfr_mode=tfr_mode,
+            zero_mean_wavelets=zero_mean_wavelets,
+            output="not_an_output",
+        )
+
     with pytest.raises(TypeError, match="`n_jobs` must be an integer."):
         compute_tfr(
             data=data,