pywt

1
import matplotlib.pyplot as plt
2
import numpy as np
3

4
import pywt
5

6

7
def gaussian(x, x0, sigma):
8
    return np.exp(-np.power((x - x0) / sigma, 2.0) / 2.0)
9

10

11
def make_chirp(t, t0, a):
12
    frequency = (a * (t + t0)) ** 2
13
    chirp = np.sin(2 * np.pi * frequency * t)
14
    return chirp, frequency
15

16

17
# generate signal
18
time = np.linspace(0, 1, 2000)
19
chirp1, frequency1 = make_chirp(time, 0.2, 9)
20
chirp2, frequency2 = make_chirp(time, 0.1, 5)
21
chirp = chirp1 + 0.6 * chirp2
22
chirp *= gaussian(time, 0.5, 0.2)
23

24
# plot signal
25
fig, axs = plt.subplots(2, 1, sharex=True)
26
axs[0].plot(time, chirp)
27
axs[1].plot(time, frequency1)
28
axs[1].plot(time, frequency2)
29
axs[1].set_yscale("log")
30
axs[1].set_xlabel("Time (s)")
31
axs[0].set_ylabel("Signal")
32
axs[1].set_ylabel("True frequency (Hz)")
33
plt.suptitle("Input signal")
34
plt.show()
35

36
# perform CWT
37
wavelet = "cmor1.5-1.0"
38
# logarithmic scale for scales, as suggested by Torrence & Compo:
39
widths = np.geomspace(1, 1024, num=100)
40
sampling_period = np.diff(time).mean()
41
cwtmatr, freqs = pywt.cwt(chirp, widths, wavelet, sampling_period=sampling_period)
42
# absolute take absolute value of complex result
43
cwtmatr = np.abs(cwtmatr[:-1, :-1])
44

45
# plot result using matplotlib's pcolormesh (image with annoted axes)
46
fig, axs = plt.subplots(2, 1)
47
pcm = axs[0].pcolormesh(time, freqs, cwtmatr)
48
axs[0].set_yscale("log")
49
axs[0].set_xlabel("Time (s)")
50
axs[0].set_ylabel("Frequency (Hz)")
51
axs[0].set_title("Continuous Wavelet Transform (Scaleogram)")
52
fig.colorbar(pcm, ax=axs[0])
53

54
# plot fourier transform for comparison
55
from numpy.fft import rfft, rfftfreq
56

57
yf = rfft(chirp)
58
xf = rfftfreq(len(chirp), sampling_period)
59
plt.semilogx(xf, np.abs(yf))
60
axs[1].set_xlabel("Frequency (Hz)")
61
axs[1].set_title("Fourier Transform")
62
plt.tight_layout()
63