pywt

1
#!/usr/bin/env python
2

3
import sys
4

5
import matplotlib.pyplot as plt
6
import numpy as np
7

8
import pywt
9

10
usage = """
11
Usage:
12
    python waveinfo.py waveletname
13

14
    Example: python waveinfo.py 'sym5'
15
"""
16

17
try:
18
    wavelet = pywt.Wavelet(sys.argv[1])
19
    try:
20
        level = int(sys.argv[2])
21
    except IndexError as e:
22
        level = 10
23
except ValueError as e:
24
    print("Unknown wavelet")
25
    raise SystemExit
26
except IndexError as e:
27
    print(usage)
28
    raise SystemExit
29

30

31
data = wavelet.wavefun(level)
32
if len(data) == 2:
33
    x = data[1]
34
    psi = data[0]
35
    fig = plt.figure()
36
    if wavelet.complex_cwt:
37
        plt.subplot(211)
38
        plt.title(wavelet.name+' real part')
39
        mi, ma = np.real(psi).min(), np.real(psi).max()
40
        margin = (ma - mi) * 0.05
41
        plt.plot(x,np.real(psi))
42
        plt.ylim(mi - margin, ma + margin)
43
        plt.xlim(x[0], x[-1])
44
        plt.subplot(212)
45
        plt.title(wavelet.name+' imag part')
46
        mi, ma = np.imag(psi).min(), np.imag(psi).max()
47
        margin = (ma - mi) * 0.05
48
        plt.plot(x,np.imag(psi))
49
        plt.ylim(mi - margin, ma + margin)
50
        plt.xlim(x[0], x[-1])
51
    else:
52
        mi, ma = psi.min(), psi.max()
53
        margin = (ma - mi) * 0.05
54
        plt.plot(x,psi)
55
        plt.title(wavelet.name)
56
        plt.ylim(mi - margin, ma + margin)
57
        plt.xlim(x[0], x[-1])
58
else:
59
    funcs, x = data[:-1], data[-1]
60
    labels = ["scaling function (phi)", "wavelet function (psi)",
61
              "r. scaling function (phi)", "r. wavelet function (psi)"]
62
    colors = ("r", "g", "r", "g")
63
    fig = plt.figure()
64
    for i, (d, label, color) in enumerate(zip(funcs, labels, colors)):
65
        mi, ma = d.min(), d.max()
66
        margin = (ma - mi) * 0.05
67
        ax = fig.add_subplot((len(data) - 1) // 2, 2, 1 + i)
68

69
        ax.plot(x, d, color)
70
        ax.set_title(label)
71
        ax.set_ylim(mi - margin, ma + margin)
72
        ax.set_xlim(x[0], x[-1])
73

74
plt.show()
75