example_wavelet_sine.py

import matplotlib.pyplot as plt 
import numpy as np 
import pandas as pd 
from waveletFunctions import wavelet, wave_signif
import pylab
from matplotlib.colors import LogNorm

"""
This is an example script of how to use the Torrence & Compo
python wavelet code. Here we'll just generate a sine wave and then also
a combination of a sine and cosine wave. 

There are two main functions that are used here which are
`wavelet` and `wave_signif` both of which are in the file `waveletFunctions.py`
and are imported above. For each of these functions, there are many additional
keyword arguments that can be passed - here we'll just use the default ones. 
"""

# Define sine wave
# -----------------------------
dt = 0.1
freq = 0.2 # i.e. a period of 5.
freq2 = 0.5
time = np.arange(1000)*dt #define time array
signal = np.sin(2*np.pi*freq*time) # define sine wave f(x) = sin(2pift)
# uncomment the below to test a signal with two periods
# signal = np.sin(2*np.pi*freq*time) + np.cos(2*np.pi*freq2*time) 

n = len(signal) # len of signal

mother = 'MORLET' # here we'll use a Morlet mother wavelet, this can be changed see notes in waveletFunctions.py
variance = np.std(signal, ddof=1) ** 2 # variance of the signal. 

"""
Compute the wavelet transform
-----------------------------
Now that we have the signal we want to compute the wavelet transform.
For this we'll use the `wavelet` function that takes arguments of the signal
and the dt - time step/sampling time.

It also takes keyword argments or optional inputs. Here, were basically only
passing the ones we really care about, and the rest will be set to the default values.
You can pass things like a different mother wavelet, a different mother wavelet parameter, 
define the different scales, and pad to get up to correct number of scales etc. To
see more about these, read the docstring above the `wavelet` function in waveletFunctions. 

Upon calling `wavelet` it returns the wavelet transform (wave), the period, scale and the 
cone-of-influence (coi). 
"""
wave, period, scale, coi = wavelet(signal, dt, mother="MORLET") # compute transform
power = (np.abs(wave)) ** 2 # wavelet power is defined as absolute value of transform squared
global_ws = (np.sum(power, axis=1) / n)  # define global wavelet - which is time-average over all times

"""
Compute the significance values
-------------------------------
To do these we'll use the `wave_signif` function. 
Here you need to pass the signal or the variance of the signal, the sampling time, 
and the scale (which is returned from above function).

In this function you can also pass other keyword arguments including the type of 
significance test you want to do (default chi-squared from paper), a `lag1` parameter
if you want to estimate red-noise, a significance level (default 0.95/95%) and the 
degrees of freedom.
"""
signif = wave_signif(variance, dt, scale, mother=mother) # calculate the significance 
sig95 = signif[:, np.newaxis].dot(np.ones(n)[np.newaxis, :])  # expand signif --> (J+1)x(N) array
sig95 = power / sig95  # where ratio > 1, power is significant

# Global wavelet spectrum & significance levels: same but for global calclation we are
# doing a time-averaged significance test - hence we pass `sigtest=1`
dof = n - scale  # the -scale corrects for padding at edges
global_signif = wave_signif(variance, dt, scale, sigtest=1, dof=dof, mother=mother)
glsig95 = global_signif[:, np.newaxis].dot(np.ones(n)[np.newaxis, :])  # expand signif --> (J+1)x(N) array



"""
Plot the results!
-------------------------------
Here we'll plot the results. I'm using pylab to define the
position of the subplot axes - but you could of course use 
just plt.subplots.

You can change anything here, this is just an example of how I plot the results.
"""
fig = pylab.figure(figsize = (11,8))
# define each position of the axes
ax = pylab.axes([0.069, 0.54, 0.68, 0.43])
bx = pylab.axes([0.069, 0.1, 0.68, 0.43], sharex=ax)
cx = pylab.axes([0.755, 0.1, 0.19, 0.43], sharey=bx)

##### First plot the input timeseries ####
ax.plot(time, signal, 'k', linewidth=2.0, color = 'black')
ax.plot(time, signal, 'k', linewidth=1.0, color = 'grey')
ax.set_xlim(time[0], time[-1])
ax.tick_params(axis="x", which="both", labelbottom=False)

#### Plot the wavelet power spectrum #####
levels = np.linspace(0, np.max((power)), 1000) # you can change these levels to highlight the power
cs = bx.contourf(time, period, power, level=levels, cmap="viridis", extend='both') # plot the power
bx.contour(time, period, glsig95, [-99, 1], colors='k', linewidths=1.) # plot the significance level contours
bx.plot(time, coi[:-1], 'white') # plot the cone of influence
bx.set_ylabel('Period (seconds) ')
bx.set_xlabel('Time (seconds) ')
bx.set_xlim(time[0], time[-1])
bx.invert_yaxis()
cb = plt.colorbar(cs, orientation = 'vertical') # colorbar for the power plot

#### Plot the global wavelet spectrum #####
cx.plot(global_signif, (period), 'k--', color = 'r', label = 'Sig 95%', lw = 1.5)
cx.plot(global_ws, period, 'k-', linewidth=1.5, color = 'k')
cx.legend(loc = 'upper right')
cb.set_label('Wavelet power')
cx.set_title('Global Wavelet')
cx.set_xlabel(r'Power')

cx.set_xscale('log')
cx.set_yscale("log")
cx.set_ylim(period.min(), period.max())
cx.tick_params(axis='y', which='major', labelleft=False)

plt.savefig("Example_wavelet.png", dpi=200)
plt.show()