pywt

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

4
import pywt
5

6
wav = pywt.ContinuousWavelet('cmor1.5-1.0')
7

8
# print the range over which the wavelet will be evaluated
9
print("Continuous wavelet will be evaluated over the range "
10
      f"[{wav.lower_bound}, {wav.upper_bound}]")
11

12
width = wav.upper_bound - wav.lower_bound
13

14
scales = [1, 2, 3, 4, 10, 15]
15

16
max_len = int(np.max(scales)*width + 1)
17
t = np.arange(max_len)
18
fig, axes = plt.subplots(len(scales), 2, figsize=(12, 6))
19
for n, scale in enumerate(scales):
20

21
    # The following code is adapted from the internals of cwt
22
    int_psi, x = pywt.integrate_wavelet(wav, precision=10)
23
    step = x[1] - x[0]
24
    j = np.floor(
25
        np.arange(scale * width + 1) / (scale * step))
26
    if np.max(j) >= np.size(int_psi):
27
        j = np.delete(j, np.where(j >= np.size(int_psi))[0])
28
    j = j.astype(np.int_)
29

30
    # normalize int_psi for easier plotting
31
    int_psi /= np.abs(int_psi).max()
32

33
    # discrete samples of the integrated wavelet
34
    filt = int_psi[j][::-1]
35

36
    # The CWT consists of convolution of filt with the signal at this scale
37
    # Here we plot this discrete convolution kernel at each scale.
38

39
    nt = len(filt)
40
    t = np.linspace(-nt//2, nt//2, nt)
41
    axes[n, 0].plot(t, filt.real, t, filt.imag)
42
    axes[n, 0].set_xlim([-max_len//2, max_len//2])
43
    axes[n, 0].set_ylim([-1, 1])
44
    axes[n, 0].text(50, 0.35, f'scale = {scale}')
45

46
    f = np.linspace(-np.pi, np.pi, max_len)
47
    filt_fft = np.fft.fftshift(np.fft.fft(filt, n=max_len))
48
    filt_fft /= np.abs(filt_fft).max()
49
    axes[n, 1].plot(f, np.abs(filt_fft)**2)
50
    axes[n, 1].set_xlim([-np.pi, np.pi])
51
    axes[n, 1].set_ylim([0, 1])
52
    axes[n, 1].set_xticks([-np.pi, 0, np.pi])
53
    axes[n, 1].set_xticklabels([r'$-\pi$', '0', r'$\pi$'])
54
    axes[n, 1].grid(True, axis='x')
55
    axes[n, 1].text(np.pi/2, 0.5, f'scale = {scale}')
56

57
axes[n, 0].set_xlabel('time (samples)')
58
axes[n, 1].set_xlabel('frequency (radians)')
59
axes[0, 0].legend(['real', 'imaginary'], loc='upper left')
60
axes[0, 1].legend(['Power'], loc='upper left')
61
axes[0, 0].set_title('filter')
62
axes[0, 1].set_title(r'|FFT(filter)|$^2$')
63