[ENH] Add option to store and return TFR taper weights #12910

tsbinns · 2024-10-22T09:27:16Z

Reference issue (if any)

What does this implement/fix?

Adds an option to return taper weights for complex and phase outputs of the multitaper method in tfr_array_multitaper(), and also ensures taper weights are stored in TFR objects.

Additional information

When working on this, I discovered a couple of other issues with the per-taper TFR implementations (#12851 (comment)), including the fact that the TFR object plotting methods and to_data_frame methods do not account for a taper dimension, leading to errors. Wasn't sure if people want me to also address these here or in a separate PR.

tsbinns · 2024-10-22T09:30:06Z

mne/time_frequency/tfr.py

@@ -302,12 +306,15 @@ def _make_dpss(
                real_offset = Wk.mean()
                Wk -= real_offset
            Wk /= np.sqrt(0.5) * np.linalg.norm(Wk.ravel())
+            Ck = np.sqrt(conc[m])


This I am somewhat unsure on. The existing implementation is to just use conc as-is, however in the MNE-Connectivity implementation that sqrt is taken: https://github.com/mne-tools/mne-connectivity/blob/97147a57eefb36a5c9680e539fdc6343a1183f20/mne_connectivity/spectral/time.py#L825

I am also unsure on this point. We should ask @ruuskas (who wrote the implementation in MNE-Connectivity) and @larsoner (who wrote the SciPy DPSS implementation) to weigh in.

I noticed for the PSD computation that the square root of the weights is also taken, so I think this is okay:

mne-python/mne/time_frequency/multitaper.py

Line 412 in b329515

weights = np.sqrt(eigvals)[np.newaxis, :, np.newaxis]

mne/time_frequency/tfr.py

tsbinns · 2024-10-22T09:36:50Z

I'm also somewhat confused about the design of the _make_dpss function:

mne-python/mne/time_frequency/tfr.py

Lines 285 to 315 in 82fc2f7

    
           for m in range(n_taps): 
        
               Wm = list() 
        
               Cm = list() 
        
               for k, f in enumerate(freqs): 
        
                   if len(n_cycles) != 1: 
        
                       this_n_cycles = n_cycles[k] 
        
                   else: 
        
                       this_n_cycles = n_cycles[0] 
        
                   t_win = this_n_cycles / float(f) 
        
                   t = np.arange(0.0, t_win, 1.0 / sfreq) 
        
                   # Making sure wavelets are centered before tapering 
        
                   oscillation = np.exp(2.0 * 1j * np.pi * f * (t - t_win / 2.0)) 
        
                   # Get dpss tapers 
        
                   tapers, conc = dpss_windows( 
        
                       t.shape[0], time_bandwidth / 2.0, n_taps, sym=False 
        
                   ) 
        
                   Wk = oscillation * tapers[m] 
        
                   if zero_mean:  # to make it zero mean 
        
                       real_offset = Wk.mean() 
        
                       Wk -= real_offset 
        
                   Wk /= np.sqrt(0.5) * np.linalg.norm(Wk.ravel()) 
        
                   Ck = np.sqrt(conc[m]) 
        
                   Wm.append(Wk) 
        
                   Cm.append(Ck) 
        
               Ws.append(Wm) 
        
               Cs.append(Cm)

It is looping over tapers, and then over frequencies. However, the dpss_windows function it calls internally provides the tapers and their weights for all tapers of a given frequency.

Would it not be more efficient to only loop over frequencies and take advantage of the fact that this will also return information for each taper?

… into add_tfr_weights

mne/time_frequency/tfr.py

This reverts commit 82fc2f7.

This reverts commit 8c16716.

tsbinns · 2024-10-22T09:58:34Z

I also have a question regarding testing: for the I/O tests, we're reading TFR objects that do not have a weights property (just gets assigned to None) when loaded. Do I need to create new TFR objects that actually have some weights, or is the current test sufficient?

Apart from this there are still some tests I need to expand.

mne/time_frequency/multitaper.py

mne/time_frequency/tfr.py

drammock · 2024-10-28T21:34:59Z

mne/time_frequency/tfr.py

@@ -302,12 +306,15 @@ def _make_dpss(
                real_offset = Wk.mean()
                Wk -= real_offset
            Wk /= np.sqrt(0.5) * np.linalg.norm(Wk.ravel())
+            Ck = np.sqrt(conc[m])


I am also unsure on this point. We should ask @ruuskas (who wrote the implementation in MNE-Connectivity) and @larsoner (who wrote the SciPy DPSS implementation) to weigh in.

mne/time_frequency/tfr.py

tsbinns · 2024-10-29T19:02:38Z

Thanks for the review @drammock! I will sort out those remaining tests, although I'm in the process of moving at the moment so it might not be for some days.

Regarding those issues I came across with TFR multitapers and converting to dataframes / plotting: would you like me to incorporate that into this PR?

drammock · 2024-12-09T20:08:35Z

Do I need to create new TFR objects that actually have some weights, or is the current test sufficient?

Yes I think we should. most (all?) of them are created by pytest fixtures at present. I see 3 options:

tweak the fixtures to always return TFRs that have weights.
when you want to test something specific to weights, monkey-patch some weights (and a taper dim) into the object at the start of the test
write a new fixture (or parametrize an existing one) so that you can get TFRs with/without weights at need.

To really test thoroughly, option (2) is probably best, because then you can also patch in things that are expected to fail, and test that they do fail in the expected way.

tsbinns · 2024-12-10T14:58:19Z

mne/time_frequency/tfr.py

@@ -1392,7 +1421,6 @@ def __setstate__(self, state):

        defaults = dict(
            method="unknown",
-            dims=("epoch", "channel", "freq", "time")[-state["data"].ndim :],


Have removed dims being set in BaseTFR since the possibility of the optional epoch and taper dimensions makes it really difficult to disentangle here. It's much easier to handle this in the individual RawTFR, EpochsTFR, and AverageTFR classes.

tsbinns · 2024-12-10T14:59:24Z

mne/time_frequency/tfr.py

+        # Set dims now since optional tapers makes it difficult to disentangle later
+        state["dims"] = ("channel",)
+        if state["data"].ndim == 4:
+            state["dims"] += ("taper",)
+        state["dims"] += ("freq", "time")


Example of handling dims in the AverageTFR class where only one dimension (taper) is optional.

tsbinns · 2024-12-10T15:04:14Z

mne/time_frequency/tfr.py

+
+        Averaging is not supported for data containing a taper dimension.
        """
+        if "taper" in self._dims:
+            raise NotImplementedError(
+                "Averaging multitaper tapers across epochs, frequencies, or times is "
+                "not supported. If averaging across epochs, consider averaging the "
+                "epochs before computing the complex/phase spectrum."
+            )


In terms of averaging for data with tapers, I went for the same approach we're using for Spectrum and just disallowing this.

I don't think this is an API change requiring a deprecation cycle since:

the docstring expects the data to not have a taper dimension, e.g. If callable, must take a NumPy array of shape (n_epochs, n_channels, n_freqs, n_times).

trying to call this method on an object with a taper dimension would raise an uncaught error: n_epochs, n_channels, n_freqs, n_times = self.data.shape (wouldn't be able to unpack this properly).

So explicitly preventing this method being called with a taper dimension doesn't change current behaviour, it just gives a nicer error as to why this can't be done.

tsbinns · 2024-12-10T16:06:53Z

mne/time_frequency/tfr.py

    Notes
    -----
+    Aggregating multitaper TFR datasets with a taper dimension such as for complex or
+    phase data is not supported.
+
    .. versionadded:: 0.11.0
    """
+    if any("taper" in tfr._dims for tfr in all_tfr):
+        raise NotImplementedError(
+            "Aggregating multitaper tapers across TFR datasets is not supported."
+        )
+


It's a similar case to averaging for the time_frequency.combine_tfr() function (which also gets called by the grand_average() function).

However, unlike the EpochsTFR.average() method, this could be considered an API change since combine_tfr() should currently run with taper data. Does preventing this use case require a deprecation cycle?

On a side note, I noticed that while a public function, combine_tfr() is not listed in the API (the equivalent combine_evoked() is). Is this an oversight or an intended omission?

… into add_tfr_weights

tsbinns · 2024-12-11T21:17:39Z

Those recent pushes added support for data with a tapers dimension in the ...TFRArray objects which was no fully accounted for before.

tsbinns · 2024-12-11T21:42:28Z

Now to_data_frame works for data with a tapers dimension (alongside unit tests).

Just sorting the issues with plotting to go!

tsbinns · 2024-12-12T12:39:19Z

Just looking into the sorting the power this morning and I am a little confused by the procedure being used to convert the complex taper coeffs into power, as it seems that no taper weights are ever applied. I opened an issue to try and figure out if this is a mistake, or a misunderstanding on my part: #13023

tsbinns

Latest push adds support for plotting of data with a taper dimension (aggregates over tapers before plotting and converts to power (if complex coeffs) or keeps as phase data).
Also has test coverage.

tsbinns · 2024-12-14T19:04:48Z

mne/time_frequency/tests/test_tfr.py

+@pytest.mark.parametrize("output", ("complex", "phase"))
+def test_tfr_topo_plotting_multitaper_complex_phase(output, evoked):
+    """Test plot_joint/topo/topomap() for data with a taper dimension."""
+    # Compute TFR with taper dimension
+    tfr = evoked.compute_tfr(
+        method="multitaper", freqs=freqs_linspace, n_cycles=4, output=output
+    )
+    # Check that plotting works
+    tfr.plot_joint(topomap_args=dict(res=8, contours=0, sensors=False))  # for speed
+    tfr.plot_topo()
+    tfr.plot_topomap()


Basic test that just checks whether the code runs, but it covers the lines where changes to topo-related plotting were made, and other tests deal with non-default method params.

tsbinns · 2024-12-14T19:06:40Z

mne/time_frequency/tests/test_tfr.py

+@pytest.mark.parametrize("output", ("complex", "phase"))
+def test_plot_multitaper_complex_phase(output):
+    """Test TFR plotting of data with a taper dimension."""
+    # Create example data with a taper dimension
+    n_chans, n_tapers, n_freqs, n_times = (3, 4, 2, 3)
+    data = np.random.rand(n_chans, n_tapers, n_freqs, n_times)
+    if output == "complex":
+        data = data + np.random.rand(*data.shape) * 1j  # add imaginary data
+    times = np.arange(n_times)
+    freqs = np.arange(n_freqs)
+    weights = np.random.rand(n_tapers, n_freqs)
+    info = mne.create_info(n_chans, 1000.0, "eeg")
+    tfr = AverageTFRArray(
+        info=info, data=data, times=times, freqs=freqs, weights=weights
+    )
+    # Check that plotting works
+    tfr.plot()


Again, a pretty basic test that just checks whether plotting code runs, but covers the changes and non-default params tested elswehere.

tsbinns · 2024-12-14T19:08:40Z

mne/time_frequency/tfr.py

-        # TODO this is the only remaining call to _preproc_tfr; should be refactored
-        #      (to use _prep_data_for_plot?)
-        data, times, freqs, vmin, vmax = _preproc_tfr(
+        # baseline, crop, convert complex to power, aggregate tapers, and dB scaling
+        data, times, freqs = _prep_data_for_plot(
            data,
            times,
            freqs,
-            tmin,
-            tmax,
-            fmin,
-            fmax,
-            mode,
-            baseline,
-            vmin,
-            vmax,
-            dB,
-            info["sfreq"],
+            tmin=tmin,
+            tmax=tmax,
+            fmin=fmin,
+            fmax=fmax,
+            baseline=baseline,
+            mode=mode,
+            dB=dB,
+            taper_weights=self.weights,
+            verbose=verbose,
        )
+        # get vlims
+        vmin, vmax = _setup_vmin_vmax(data, vmin, vmax)


Seemed like as good as time as any to refactor and replace the _preproc_tfr call with _prep_data_for_plot where changes for handling data with a taper dimension have been made.

tsbinns · 2024-12-14T19:10:56Z

mne/viz/topomap.py

+    # handle unaggregated multitaper (complex or phase multitaper data)
+    if tfr.weights is not None:  # assumes a taper dimension
+        logger.info("Aggregating multitaper estimates before plotting...")
+        weights = tfr.weights[np.newaxis, :, :, np.newaxis]  # add channel & time dims
+        data = weights * data
+        if np.iscomplexobj(data):  # complex coefficients → power
+            data *= data.conj()
+            data = data.real.sum(axis=1)
+            data *= 2 / (weights * weights.conj()).real.sum(axis=1)
+        else:  # tapered phase data → weighted phase data
+            data = data.mean(axis=1)


Also the case that viz.plot_tfr_topomap() needs to be able to handle data with a tapers dim. Unfortunately circular imports mean the code for handling taper dims from tfr.py can't be used here, so there's a bit of code repetition.

tsbinns · 2024-12-14T19:12:34Z

mne/time_frequency/tfr.py

+        else:  # tapered phase data → weighted phase data
+            data = (data * taper_weights[np.newaxis, :, :, np.newaxis]).mean(axis=1)


This is my guess at aggregating over tapers for phase data. Can anyone confirm if this is correct?

tsbinns · 2024-12-14T19:15:38Z

mne/time_frequency/tfr.py

+    tfr = weights * x_mt
+    tfr *= tfr.conj()
+    tfr = tfr.real.sum(axis=1)
+    tfr *= 2 / (weights * weights.conj()).real.sum(axis=1)


This aggregation over tapers for the complex spectra follows the same procedure for computing PSDs and how we handle TFR data in MNE-Connectivity.
However, it does differ to how this aggregation is handled elsewhere in the TFR classes (see #13023), so it would be nice to clarify the correct approach.

tsbinns · 2024-12-19T12:31:45Z

Just looking into the sorting the power this morning and I am a little confused by the procedure being used to convert the complex taper coeffs into power, as it seems that no taper weights are ever applied. I opened an issue to try and figure out if this is a mistake, or a misunderstanding on my part: #13023

As discussed there this is a bug but will be addressed in a separate PR once this is merged.

tsbinns added 3 commits October 22, 2024 11:17

Add option to store and return tfr taper weights

9fe1fb6

Merge remote-tracking branch 'upstream/main' into add_tfr_weights

45c6a0b

Update docstrings

82fc2f7

tsbinns requested review from drammock, adam2392 and mscheltienne as code owners October 22, 2024 09:27

tsbinns commented Oct 22, 2024

View reviewed changes

mne/time_frequency/tfr.py Show resolved Hide resolved

tsbinns added 3 commits October 22, 2024 11:36

Merge branch 'main' into add_tfr_weights

9f30a59

Remove whitespace

a49f934

Merge branch 'add_tfr_weights' of https://github.com/tsbinns/mne-python…

48afced

… into add_tfr_weights

tsbinns commented Oct 22, 2024

View reviewed changes

mne/time_frequency/tfr.py Show resolved Hide resolved

tsbinns added 5 commits October 22, 2024 11:39

Add PR num

7c3dcfa

Revert "Update docstrings"

8c16716

This reverts commit 82fc2f7.

Remove outdated default setting

51b8cd0

Reapply "Update docstrings"

2f9a4b4

This reverts commit 8c16716.

Update docstrings

b4537b2

tsbinns added 2 commits October 24, 2024 19:30

Merge branch 'main' into add_tfr_weights

f155238

Merge branch 'main' into add_tfr_weights

2a03e9b

drammock reviewed Oct 28, 2024

View reviewed changes

tsbinns added 2 commits October 29, 2024 19:57

Merge branch 'main' into add_tfr_weights

045d9a2

Enforce return_weights as named param

8d645bb

tsbinns added 3 commits December 9, 2024 10:17

Merge branch 'main' into add_tfr_weights

5ad9bd5

Add missing test coverage

1c02b40

Add changelog entry

54f2a32

tsbinns requested review from larsoner and agramfort as code owners December 9, 2024 11:43

Fix docstring entries

ca27179

tsbinns mentioned this pull request Dec 10, 2024

Store n_cycles and time_bandwidth params in *TFR objects #12851

Open

tsbinns added 4 commits December 10, 2024 13:45

Fix faulty state check

b14a100

Add weights to AverageTFR

972aba2

Expand test coverage

e11fa2b

Merge branch 'main' into add_tfr_weights

aaef4b7

tsbinns commented Dec 10, 2024

View reviewed changes

Disallow aggregating tapers in combine_tfr

999d122

tsbinns commented Dec 10, 2024

View reviewed changes

tsbinns added 6 commits December 10, 2024 16:22

Updated docstrings

e12b09a

Merge branch 'main' into add_tfr_weights

dd61955

Add placeholder versionadded tags

728701e

Merge branch 'add_tfr_weights' of https://github.com/tsbinns/mne-python…

6af3310

… into add_tfr_weights

Merge remote-tracking branch 'upstream/main' into add_tfr_weights

e3a3c4b

Begin fixing to_data_frame

de39d25

Fix to_data_frame bug with tapers

80126a7

tsbinns mentioned this pull request Dec 12, 2024

Possible bug in computation of multitaper TFR power #13023

Open

Fix plotting with tapers

82dfab9

tsbinns commented Dec 14, 2024

View reviewed changes

Merge branch 'main' into add_tfr_weights

5b150aa

larsoner added this to the 1.10 milestone Dec 16, 2024

Merge branch 'main' into add_tfr_weights

0d3d85d

Add version tag

012bd94

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ENH] Add option to store and return TFR taper weights #12910

[ENH] Add option to store and return TFR taper weights #12910

tsbinns commented Oct 22, 2024

tsbinns Oct 22, 2024

drammock Oct 28, 2024

tsbinns Dec 12, 2024

tsbinns commented Oct 22, 2024

tsbinns commented Oct 22, 2024 •

edited

Loading

drammock Oct 28, 2024

tsbinns commented Oct 29, 2024 •

edited

Loading

drammock commented Dec 9, 2024

tsbinns Dec 10, 2024

tsbinns Dec 10, 2024

tsbinns Dec 10, 2024 •

edited

Loading

tsbinns Dec 10, 2024 •

edited

Loading

tsbinns commented Dec 11, 2024

tsbinns commented Dec 11, 2024 •

edited

Loading

tsbinns commented Dec 12, 2024

tsbinns left a comment

tsbinns Dec 14, 2024

tsbinns Dec 14, 2024

tsbinns Dec 14, 2024

tsbinns Dec 14, 2024 •

edited

Loading

tsbinns Dec 14, 2024

tsbinns Dec 14, 2024

tsbinns commented Dec 19, 2024

		else: # tapered phase data → weighted phase data
		data = (data * taper_weights[np.newaxis, :, :, np.newaxis]).mean(axis=1)

[ENH] Add option to store and return TFR taper weights #12910

Are you sure you want to change the base?

[ENH] Add option to store and return TFR taper weights #12910

Conversation

tsbinns commented Oct 22, 2024

Reference issue (if any)

What does this implement/fix?

Additional information

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tsbinns commented Oct 22, 2024

tsbinns commented Oct 22, 2024 • edited Loading

Choose a reason for hiding this comment

tsbinns commented Oct 29, 2024 • edited Loading

drammock commented Dec 9, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tsbinns Dec 10, 2024 • edited Loading

Choose a reason for hiding this comment

tsbinns Dec 10, 2024 • edited Loading

Choose a reason for hiding this comment

tsbinns commented Dec 11, 2024

tsbinns commented Dec 11, 2024 • edited Loading

tsbinns commented Dec 12, 2024

tsbinns left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tsbinns Dec 14, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tsbinns commented Dec 19, 2024

tsbinns commented Oct 22, 2024 •

edited

Loading

tsbinns commented Oct 29, 2024 •

edited

Loading

tsbinns Dec 10, 2024 •

edited

Loading

tsbinns Dec 10, 2024 •

edited

Loading

tsbinns commented Dec 11, 2024 •

edited

Loading

tsbinns Dec 14, 2024 •

edited

Loading