pywt

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

4
import pywt
5

6
# plot complex morlet wavelets with different center frequencies and bandwidths
7
wavelets = [f"cmor{x:.1f}-{y:.1f}" for x in [0.5, 1.5, 2.5] for y in [0.5, 1.0, 1.5]]
8
fig, axs = plt.subplots(3, 3, figsize=(10, 10), sharex=True, sharey=True)
9
for ax, wavelet in zip(axs.flatten(), wavelets):
10
    [psi, x] = pywt.ContinuousWavelet(wavelet).wavefun(10)
11
    ax.plot(x, np.real(psi), label="real")
12
    ax.plot(x, np.imag(psi), label="imag")
13
    ax.set_title(wavelet)
14
    ax.set_xlim([-5, 5])
15
    ax.set_ylim([-0.8, 1])
16
ax.legend()
17
plt.suptitle("Complex Morlet Wavelets with different center frequencies and bandwidths")
18
plt.show()
19

20

21
def gaussian(x, x0, sigma):
22
    return np.exp(-np.power((x - x0) / sigma, 2.0) / 2.0)
23

24

25
def make_chirp(t, t0, a):
26
    frequency = (a * (t + t0)) ** 2
27
    chirp = np.sin(2 * np.pi * frequency * t)
28
    return chirp, frequency
29

30

31
def plot_wavelet(time, data, wavelet, title, ax):
32
    widths = np.geomspace(1, 1024, num=75)
33
    cwtmatr, freqs = pywt.cwt(
34
        data, widths, wavelet, sampling_period=np.diff(time).mean()
35
    )
36
    cwtmatr = np.abs(cwtmatr[:-1, :-1])
37
    pcm = ax.pcolormesh(time, freqs, cwtmatr)
38
    ax.set_yscale("log")
39
    ax.set_xlabel("Time (s)")
40
    ax.set_ylabel("Frequency (Hz)")
41
    ax.set_title(title)
42
    plt.colorbar(pcm, ax=ax)
43
    return ax
44

45

46
# generate signal
47
time = np.linspace(0, 1, 1000)
48
chirp1, frequency1 = make_chirp(time, 0.2, 9)
49
chirp2, frequency2 = make_chirp(time, 0.1, 5)
50
chirp = chirp1 + 0.6 * chirp2
51
chirp *= gaussian(time, 0.5, 0.2)
52

53
# perform CWT with different wavelets on same signal and plot results
54
wavelets = [f"cmor{x:.1f}-{y:.1f}" for x in [0.5, 1.5, 2.5] for y in [0.5, 1.0, 1.5]]
55
fig, axs = plt.subplots(3, 3, figsize=(10, 10), sharex=True)
56
for ax, wavelet in zip(axs.flatten(), wavelets):
57
    plot_wavelet(time, chirp, wavelet, wavelet, ax)
58
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
59
plt.suptitle("Scaleograms of the same signal with different wavelets")
60
plt.show()
61