updating TFR classes

[drammock]

WIP. The goal here is to copy the new Spectrum API into the time-frequency analysis code. Briefly this means (1) a new class RawTFR, (2) new instance methods [raw,epochs,evoked].compute_tfr(), (3) legacy-ify old API entry points tfr_morlet, tfr_multitaper, and tfr_stockwell (the tfr_array_* functions will remain), and (4) flesh out / make more consistent the API of the *TFR classes.

• sketch API of compute_tfr()
• make new RawTFR class
• refactor common actions (arg checking etc) into base class __init__
• use class __init__s in the compute_tfr() methods
• deprecate / legacy-ify tfr_* functions
• update tests

[mscheltienne]

I have a couple of questions:

• where do the AverageTFR and EpochsTFR class fit in?
• what was the reason that raw objects were excluded from the functions tfr_?

[drammock]

• where do the AverageTFR and EpochsTFR class fit in?

They will be retained, possibly with some signature changes in their __init__. The WIP file mne/time-frequency/spectrogram.py will go away and be merged into mne/time-frequency/tfr.py, and e.g. BaseSpectrogram will become BaseTFR. For now it's separate just for easier WIP development.

• what was the reason that raw objects were excluded from the functions tfr_?

IIUC, because it's computationally expensive and wasn't viewed as a desired/useful use case. But in recent years a few folks have asked for it, so we're adding it.

[mscheltienne]

I briefly ran the new methods on the sample dataset, and it was very intuitive! I'm looking forward to using this new API! A couple of nitpick/changes which could improve the documentation clarity. Thank you @drammock!

[drammock]

TODO: changed examples

• Intro overview tutorial: PR vs main: now uses a one-sided colormap. This is intentional/expected, but I want to make sure folks are OK with it. (one-sided vs two-sided cmap for TFR image plots is decided empirically at plotting time, but generally speaking we should only get one-sided if no baselining was done)
• clustering results: PR vs main was a zero_mean difference
• joint plot: PR vs main the colormap limits change slightly on this PR. I view this as a bugfix; it is now guaranteed to reflect the plotted min/max across the joint topomaps (unless user overrides vlim in topomap_args).
• TFR image plot PR vs main: for n_cycles=1 the image changes slightly. needs further investigation. was a zero_mean difference
• hilbert TFR example PR vs main: something definitely wrong with data or colormap. needs further investigation. previously it was setting vlim=(-0.1, 0.1) and the colormap was implicitly being set to RdBu_r (which doesn't make sense for all-positive power data). But we do need to fix vmax at least, in order for the subplots to make sense, so I've changed it to vlim=(0, 0.1) because the code is now smarter about which cmap it chooses (and by default it will choose Reds now for this data)
• the example of "operating on arrays" should probably demo a TFRArray class rather than the manual use of centers_to_edges and pcolormesh. Done.
• single-trial induced power PR vs main needs further investigation. was a zero_mean difference

[drammock]

some of the changed examples mentioned above ☝🏻 are caused by the fact that

• the old tutorial code used tfr_morlet(epochs, ...) which defaults to zero_mean=True, use_fft=False
• the new code uses epochs.compute_tfr("morlet", ...) which ultimately wraps to tfr_array_morlet(), which has the opposite defaults (zero_mean=False, use_fft=True)

That discovery makes it easy enough to preserve the tutorial plots unchanged... but it seems really weird to have opposite defaults for tfr_morlet vs tfr_array_morlet. @agramfort or @larsoner any insight into whether that was intentional / which one should be the default going forward?

[larsoner]

That seems like it's likely a bug but we should maybe fix via deprecation. IIRC zero_mean=True is the better default since it avoids problems with the DC component in a signal leaking into higher frequencies. use_fft shouldn't make a difference in the results/numbers obtained just the speed, and I think which is faster might be data dependent, so we can probably leave that alone.

[drammock]

digging deeper, I found this docstring on the private internal _compute_tfr function:

zero_mean : bool | None, default None
    None means True for method='multitaper' and False for method='morlet'.
    If True, make sure the wavelets have a mean of zero.

Which looks pretty intentional... but still, the function (tfr_morlet()) that took MNE objects was defaulting to True and that is what I expected it to be for the corresponding array function. So I think I'll move forward with a change (with deprecation) to the default value of the array function.

[drammock]

OK I think I've resolved all the non-matching tutorial plots now! Note that in 5ea9cc7 I also did what I think is a simplification/improvement to the Hilbert example in examples/time_frequency/time_frequency_simulated.py:

• previously it would allocate a complex array (chs, freqs, times), and for each freq, put in the hilbert data after multiplying by its .conj() and taking the mean across epochs. But by not also taking the .real part, the data remained complex, which meant it was getting multiplied by its .conj() again during the plotting step. Pretty sure that was bug-ish.
• now, we don't manually multiply by .conj() (because it will still happen automatically before plotting the data), and keep the epochs separate and do the averaging all at once with the average() method after creating the EpochsTFRArray object. This latter change will necessitate more memory usage, but feels more natural / better illustrates what I think users are likely to do.

@larsoner I'd appreciate a look at just that commit to make sure you agree with that change. The resulting plot does change a little (PR vs main) but in a way that I think makes it more accurate and realistic.

[drammock]

failure on ubuntu minimal is unrelated: a GitHub timeout when checking for latest release in the sys_info test: https://github.com/mne-tools/mne-python/actions/runs/8361492634/job/22889739979?pr=11282#step:17:3485

[larsoner]

Checked the example and it looks okay, restarted the problematic CI run and marked for merge-when-green, thanks in advance @drammock !