pywt

1
#!/usr/bin/env python
2

3
import matplotlib.pyplot as plt
4
import numpy as np
5

6
import pywt
7

8
time, sst = pywt.data.nino()
9
dt = time[1] - time[0]
10

11
# Taken from http://nicolasfauchereau.github.io/climatecode/posts/wavelet-analysis-in-python/
12
wavelet = 'cmor1.5-1.0'
13
scales = np.arange(1, 128)
14

15
[cfs, frequencies] = pywt.cwt(sst, scales, wavelet, dt)
16
power = (abs(cfs)) ** 2
17

18
period = 1. / frequencies
19
levels = [0.0625, 0.125, 0.25, 0.5, 1, 2, 4, 8]
20
f, ax = plt.subplots(figsize=(15, 10))
21
ax.contourf(time, np.log2(period), np.log2(power), np.log2(levels),
22
            extend='both')
23

24
ax.set_title(f'Nino1+2 Wavelet Power Spectrum ({wavelet})')
25
ax.set_ylabel('Period (years)')
26
Yticks = 2 ** np.arange(np.ceil(np.log2(period.min())),
27
                        np.ceil(np.log2(period.max())))
28
ax.set_yticks(np.log2(Yticks))
29
ax.set_yticklabels(Yticks)
30
ax.invert_yaxis()
31
ylim = ax.get_ylim()
32
ax.set_ylim(ylim[0], -1)
33

34
plt.show()
35