Repository: neurodsp-tools/neurodsp (language: python)

Issue:
[DOC] - Tutorial for morlet wavelets
This is a tutorial for performing time-frequency decomposition using morlet wavelets following the neurodsp.timefrequency.wavelets module.


Discussion / hints:
To provide a general comment about the goal of the tutorial, I think it might help to think about what someone who wants to use the `neurodsp` morlet wavelets functionality would want to get out of a tutorial. I think they would want to know:

- How to use the functions
- That they work by recapitulating simulated data
- How and when to change the parameters

The first two goals are pretty close but I think need more explaining and I think the third could really use a lot of explaining and some plots.

@agquinn, when you open a PR you want to reference the issue that you are addressing like so:
"addresses https://github.com/neurodsp-tools/neurodsp/issues/166"

@TomDonoghue et al. @agquinn and I are working on this tutorial together as a bit of a learning exercise so feel free to review whenever or at the end I can tag you to let you know my review is done. 


 
Also morlet wavelets are a time-frequency method and the plots I saw from running the tutorial were homogenous in time so that is something that could due to be changed as well
Hi @alexrockhill and @agquinn, thanks for the PR! Feel free to tag myself (and Tom) for a review when ready.
It's a bit of a sphinx idiosyncrasy but in order for a tutorial to plot, the file name must begin with `plot_` so it should be changed to `plot_wavelet_tutorial.py`
And probably wavelet and not wavelets in the name
Hey @agquinn and @alexrockhill - this is super cool that you are doing this, a big thank you for choosing to contribute to NDSP! Adding a wavelet tutorial sounds great to me!

Please feel free to tag myself or Ryan in if you need help with anything with the module, or if it's ready for a general review. 

I've had a quick scan - for some overview notes, I'd just add that you can check the other tutorials to match conventions / layouts and that there might want to be more code / examples / plots (as compared to the amount of text), including having time varying oscillations (using `sim_bursty_oscillation`, perhaps), and exploring different settings, etc. 
 @alexrockhill I know I haven't resolved all of the issues yet, but I'm kind of stuck on how to plot the Morlet wavelet transformed series. I can't figure out why they don't look like individual bursts, and I've tried changing the frequency parameters, and messing with the normalization methods, but I'm still kind of stuck on this part. I don't know if you have any advice/ideas for this part, but I'd appreciate the help.
> @alexrockhill I know I haven't resolved all of the issues yet, but I'm kind of stuck on how to plot the Morlet wavelet transformed series. I can't figure out why they don't look like individual bursts, and I've tried changing the frequency parameters, and messing with the normalization methods, but I'm still kind of stuck on this part. I don't know if you have any advice/ideas for this part, but I'd appreciate the help.

So it's hard to review because the python file is called `wavelets_tutorial.py` and not `plot_wavelets_tutorial.py` so it doesn't run in the documentation... I know you want feedback but you really have to incorporate the feedback already given first.
It's also hard to review because you have a lot of code written and especially a ton of comments that in total is way too much. You can save the extra code somewhere outside the project directory (or on a different branch i.e. `git checkout -b tmp` and then `git checkout wavelet_tutorial` and then delete most of the wavelets_tutorial.py file but just save the good parts and then if you ever need anything you deleted you can do `git checkout tmp` and the file will be how you left it and you can copy it and `git checkout wavelet_tutorial` and paste it back in).

As for the core of the code here is mostly what you wrote with a few changes:

```
import numpy as np
import matplotlib.pyplot as plt
from neurodsp.sim import sim_bursty_oscillation
from neurodsp.utils import set_random_seed, create_times
from neurodsp.plts.time_series import plot_time_series
from neurodsp.timefrequency.wavelets import compute_wavelet_transform

set_random_seed(0)

fs = 500
n_seconds = 10
freq = 20

# Simulate a signal with bursty oscillations at 20 Hz with a decay time of 2 seconds
sig = sim_bursty_oscillation(n_seconds, fs, freq)
times = create_times(n_seconds, fs)

freqs = np.linspace(5, 100, 50)
mwt = compute_wavelet_transform(sig, fs=fs, n_cycles=7, freqs=freqs)

fig, ax = plt.subplots()
ax.imshow(abs(mwt), aspect='auto')
ax.invert_yaxis()
ax.set_xlabel('time (s)')
ax.set_xticks(np.linspace(0, times.size, 5))
ax.set_xticklabels(np.round(np.linspace(times[0], times[-1], 5), 2))
ax.set_ylabel('freq (Hz)')
ax.set_yticks(np.linspace(0, freqs.size, 5))
ax.set_yticklabels(np.round(np.linspace(freqs[0], freqs[-1], 5), 2))
fig.show()
```

Going forward, I would suggest deleting everything you have and starting with this code and adding comments to explain it. You just have so much going on in what you started that it's not very helpful for someone trying to learn about the Morlet wavelet convolution method. I would try thinking about what would you want to read if you had never heard of this so that you could code it up in python, use it and understand it. You shouldn't just copy it from here https://mne.tools/stable/auto_examples/time_frequency/plot_time_frequency_simulated.html?highlight=morlet but most of the information is already described elsewhere, the point is to make it easy for new neuroDSP users to use and understand Morlet wavelets.
That's a good start! Now it needs a couple of sentences explaining what's going on at the top and then maybe going into what's going on under the hood by making the wavelet and plotting it as you were trying to do at the end of the very long version that you had previously.

Also "Wavelets" is not as informative as "Morlet Wavelet Analysis" as a title.
@alexrockhill I think I've resolved most of the comments that came up earlier. 
This is a great start, not too text heavy now, I can read it easily and follow along. I made a few minor comments about the style etc. that would be great if you could fix.

The main thing that would be good to add, is that although this tutorial tells you how to use the Morlet Wavelet transform, it tells you nothing about what's happening under the hood. Quoting from here https://mne.tools/dev/auto_examples/time_frequency/plot_time_frequency_simulated.html "[Morlet Wavelets] are a sinusoidal wave with a gaussian envelope. We can control the balance between spectral and temporal resolution with the `n_cycles` parameter, which defines the number of cycles to include in the window." This is a good start but that's from a tutorial that compares a bunch of time frequency methods. When just describing the Morlet method it would be good to go into what exactly is happening. I would add some information from other sources perhaps starting with the wiki and finding an article to cite (https://en.wikipedia.org/wiki/Morlet_wavelet) (you'll want to look and see how the other tutorials cite articles when citing something in the tutorial). The amount of depth you want to go into it is up to you, maybe just adding some similar plots to the wiki at a minimum and reading and really going into plotting and visualizing if you have the time, but something under the hood would be helpful to users.

This might also help https://www.youtube.com/watch?v=7ahrcB5HL0k, if you can condense this video into one or two plots that would be perfect.
Also it looks like someone has already added a tutorial on wavelets...

If you want to add to it you can do

```
git checkout master
git pull upstream master
git checkout wavelet_tutorial
git rebase -i master
```

Then look at the file and fix all the conflicts between their tutorial and yours they will be separated by a line of  `=`s, then do

```
git add -A
git rebase --continue
```

Then the next push will have to be `git push origin wavelet_tutorials --force`
@alexrockhill - I'm not sure about you're previous comment: there is no current wavelet tutorial? I don't think anything should need rebasing here. 

@ryanhammonds - thanks for reviewing! Let me know if anything needs me to look over it. 
@ryanhammonds @TomDonoghue Thank you both for looking over this! I made most of the changes now, and I think it looks okay. I wasn't sure how much detail to go into, especially for the convolution part, but I think it seems pretty easy to follow. Let me know if there is anything else I should add/change!
@alexrockhill I think this looks pretty good, let me know if you have any suggestions on changing/modifying anything too!
> @alexrockhill - I'm not sure about you're previous comment: there is no current wavelet tutorial? I don't think anything should need rebasing here.
> 
> @ryanhammonds - thanks for reviewing! Let me know if anything needs me to look over it.

@TomDonoghue Audrey and I had a Zoom call and I usually never modify the master but I wasn't thinking about it at the time and then when I build the documentation the first thing I jumped to was someone already wrote a tutorial hahaha. Oops!
[Morlet Wavelet Analysis — neurodsp 2.2.0-dev documentation.pdf](https://github.com/neurodsp-tools/neurodsp/files/5850955/Morlet.Wavelet.Analysis.neurodsp.2.2.0-dev.documentation.pdf)

Looks like there is an error...
Yeah, I'm not sure why that is...I'll look into it
looks like this is super close! @ryanhammonds - could you do a sweep, after your prior review, and check for any remaining suggestions, and resolve any open comments above that have been addressed. Once your happy with it, tag me for review, and I'll have a final read / check and merge when it's ready. 
Sorry about that, I left a bad review comment that lead to the error. I pushed a fix along with a few formatting changes.
Ready to merge?
Hey all, sorry for a massive delay on my part here - I've been moving job & city etc and this fell through the cracks. 

@agquinn & @alexrockhill : a huge thank you for the submission here to improve our documentation, we greatly appreciate it! I did a full sweep of the tutorial, and because it had been so long and I didn't want to revive this with a ToDo list since y'all have probably moved on, I made direct edits to tweak some things, and add a bit of information. I'm really liking where this tutorial is now. If you have any comments or questions on the updated version, please let us know, but otherwise I don't think anything else is needed, and this should be ready to merge. 

@ryanhammonds - I think this should be good now. Can you do a quick sweep and check I didn't introduce any weird typos / bugs, and then merge if it looks good to you?

Fix patch:
diff --git a/doc/api.rst b/doc/api.rst
index 0659ec13..1ededf70 100644
--- a/doc/api.rst
+++ b/doc/api.rst
@@ -325,3 +325,12 @@ Rhythm
 
     plot_swm_pattern
     plot_lagged_coherence
+
+Time Frequency
+~~~~~~~~~~~~~~
+
+.. currentmodule:: neurodsp.plts
+.. autosummary::
+    :toctree: generated/
+
+    plot_timefrequency
diff --git a/neurodsp/plts/__init__.py b/neurodsp/plts/__init__.py
index 4abd04b3..f681b75a 100644
--- a/neurodsp/plts/__init__.py
+++ b/neurodsp/plts/__init__.py
@@ -5,3 +5,4 @@
 from .rhythm import plot_swm_pattern, plot_lagged_coherence
 from .spectral import (plot_power_spectra, plot_spectral_hist,
                        plot_scv, plot_scv_rs_lines, plot_scv_rs_matrix)
+from .timefrequency import plot_timefrequency
diff --git a/neurodsp/plts/timefrequency.py b/neurodsp/plts/timefrequency.py
new file mode 100644
index 00000000..f1d64208
--- /dev/null
+++ b/neurodsp/plts/timefrequency.py
@@ -0,0 +1,72 @@
+"""Plotting functions for neurodsp.timefrequency."""
+
+import numpy as np
+
+from neurodsp.plts.style import style_plot
+from neurodsp.plts.utils import check_ax, savefig
+
+###################################################################################################
+###################################################################################################
+
+@savefig
+@style_plot
+def plot_timefrequency(times, freqs, powers, x_ticks=5, y_ticks=5, ax=None, **kwargs):
+    """Plot a time-frequency representation of data.
+
+    Parameters
+    ----------
+    times : 1d array
+        The time dimension for the time-frequency representation.
+    freqs : 1d array
+        The frequency dimension for the time-frequency representation.
+    powers : 2d array
+        Power values to plot.
+        If array is complex, the real component is taken for plotting.
+    x_ticks, y_ticks : int or array_like
+        Defines the tick labels to add to the plot.
+        If int, is the number of evenly sampled labels to add to the plot.
+        If array_like, is a set of labels to add to the plot.
+    ax : matplotlib.Axes, optional
+        Figure axes upon which to plot.
+    **kwargs
+        Keyword arguments for customizing the plot.
+
+    Examples
+    --------
+    Plot a Morlet transformation:
+
+    >>> import numpy as np
+    >>> from neurodsp.sim import sim_bursty_oscillation
+    >>> from neurodsp.timefrequency.wavelets import compute_wavelet_transform
+    >>> fs=1000
+    >>> sig = sim_bursty_oscillation(n_seconds=10, fs=fs, freq=10)
+    >>> times = np.arange(0, len(sig)/fs, 1/fs)
+    >>> freqs = np.arange(1, 50, 1)
+    >>> mwt = compute_wavelet_transform(sig, fs, freqs)
+    >>> plot_timefrequency(times, freqs, mwt)
+    """
+
+    ax = check_ax(ax, None)
+
+    if np.iscomplexobj(powers):
+        powers = abs(powers)
+
+    ax.imshow(powers, aspect='auto', **kwargs)
+    ax.invert_yaxis()
+
+    ax.set_xlabel('Time (s)')
+    ax.set_ylabel('Frequency (Hz)')
+
+    if isinstance(x_ticks, int):
+        x_tick_pos = np.linspace(0, times.size, x_ticks)
+        x_ticks = np.round(np.linspace(times[0], times[-1], x_ticks), 2)
+    else:
+        x_tick_pos = [np.argmin(np.abs(times - val)) for val in x_ticks]
+    ax.set(xticks=x_tick_pos, xticklabels=x_ticks)
+
+    if isinstance(y_ticks, int):
+        y_ticks_pos = np.linspace(0, freqs.size, y_ticks)
+        y_ticks = np.round(np.linspace(freqs[0], freqs[-1], y_ticks), 2)
+    else:
+        y_ticks_pos = [np.argmin(np.abs(freqs - val)) for val in y_ticks]
+    ax.set(yticks=y_ticks_pos, yticklabels=y_ticks)
diff --git a/tutorials/timefreq/plot_MorletWavelet.py b/tutorials/timefreq/plot_MorletWavelet.py
index 9e9db717..8a50542f 100644
--- a/tutorials/timefreq/plot_MorletWavelet.py
+++ b/tutorials/timefreq/plot_MorletWavelet.py
@@ -18,7 +18,7 @@
 
 # Import simulation and plot code to create & visualize data
 from neurodsp.sim import sim_combined
-from neurodsp.plts import plot_time_series
+from neurodsp.plts import plot_time_series, plot_timefrequency
 from neurodsp.utils import create_times
 
 # Import function for Morlet Wavelets
@@ -85,13 +85,7 @@
 ###################################################################################################
 
 # Plot morlet wavelet transform
-_, ax = plt.subplots()
-ax.imshow(abs(mwt), aspect='auto')
-ax.invert_yaxis()
-ax.set(xlabel='Time (s)', ylabel='Frequency (Hz)',
-       xticks=np.linspace(0, times.size, 5), yticks=np.linspace(0, freqs.size, 5),
-       xticklabels=np.round(np.linspace(times[0], times[-1], 5), 2),
-       yticklabels=np.round(np.linspace(freqs[0], freqs[-1], 5), 2))
+plot_timefrequency(times, freqs, mwt)
 
 ###################################################################################################
 #
@@ -124,13 +118,8 @@
 ###################################################################################################
 
 # Plot morlet wavelet transform
-_, ax = plt.subplots()
-ax.imshow(abs(mwt), aspect='auto')
-ax.invert_yaxis()
-ax.set(xlabel='Time (s)', ylabel='Frequency (Hz)',
-       xticks=np.linspace(0, times.size, 5), yticks=np.linspace(0, freqs.size, 5),
-       xticklabels=np.round(np.linspace(times[0], times[-1], 5), 2),
-       yticklabels=np.round(np.linspace(freqs[0], freqs[-1], 5), 2))
+plot_timefrequency(times, freqs, mwt)
+
 
 ###################################################################################################
 #
@@ -279,13 +268,7 @@
 mwt = compute_wavelet_transform(sig, fs=fs, n_cycles=15, freqs=freqs)
 
 # Plot the wavelet transform
-_, ax = plt.subplots()
-ax.imshow(abs(mwt), aspect='auto')
-ax.invert_yaxis()
-ax.set(xlabel='Time (s)', ylabel='Frequency (Hz)',
-       xticks=np.linspace(0, times.size, 5), yticks=np.linspace(0, freqs.size, 5),
-       xticklabels=np.round(np.linspace(times[0], times[-1], 5), 2),
-       yticklabels=np.round(np.linspace(freqs[0], freqs[-1], 5), 2))
+plot_timefrequency(times, freqs, mwt)
 
 ###################################################################################################
 #
@@ -308,13 +291,7 @@
 mwt = compute_wavelet_transform(sig, fs=fs, n_cycles=15, freqs=freqs)
 
 # Plot the wavelet transform
-fig, ax = plt.subplots()
-ax.imshow(abs(mwt), aspect='auto')
-ax.invert_yaxis()
-ax.set(xlabel='Time (s)', ylabel='Frequency (Hz)',
-       xticks=np.linspace(0, times.size, 5), yticks=np.linspace(0, freqs.size, 5),
-       xticklabels=np.round(np.linspace(times[0], times[-1], 5), 2),
-       yticklabels=np.round(np.linspace(freqs[0], freqs[-1], 5), 2))
+plot_timefrequency(times, freqs, mwt)
 
 ###################################################################################################
 #


Test patch:
diff --git a/neurodsp/tests/plts/test_timefrequency.py b/neurodsp/tests/plts/test_timefrequency.py
new file mode 100644
index 00000000..b1e703cd
--- /dev/null
+++ b/neurodsp/tests/plts/test_timefrequency.py
@@ -0,0 +1,29 @@
+"""Tests for neurodsp.plts.timefrequency."""
+
+import numpy as np
+
+from neurodsp.utils import create_times
+from neurodsp.timefrequency.wavelets import compute_wavelet_transform
+
+from neurodsp.tests.settings import TEST_PLOTS_PATH, N_SECONDS, FS
+from neurodsp.tests.tutils import plot_test
+
+from neurodsp.plts.timefrequency import *
+
+###################################################################################################
+###################################################################################################
+
+@plot_test
+def test_plot_timefrequency(tsig_comb):
+
+    times = create_times(N_SECONDS, FS)
+    freqs = np.array([5, 10, 15, 20, 25])
+    mwt = compute_wavelet_transform(tsig_comb, FS, freqs)
+
+    plot_timefrequency(times, freqs, mwt,
+                       save_fig=True, file_path=TEST_PLOTS_PATH,
+                       file_name='test_plot_timefrequency1.png')
+
+    plot_timefrequency(times, freqs, mwt, x_ticks=[2.5, 7.5], y_ticks=[5, 15],
+                       save_fig=True, file_path=TEST_PLOTS_PATH,
+                       file_name='test_plot_timefrequency2.png')