pywt

1
#!/usr/bin/env python
2

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

6
import pywt.data
7
from pywt import WaveletPacket2D
8

9
arr = pywt.data.aero()
10

11
wp2 = WaveletPacket2D(arr, 'db2', 'symmetric', maxlevel=2)
12

13
# Show original figure
14
plt.imshow(arr, interpolation="nearest", cmap=plt.cm.gray)
15

16
path = ['d', 'v', 'h', 'a']
17

18
# Show level 1 nodes
19
fig = plt.figure()
20
for i, p2 in enumerate(path):
21
    ax = fig.add_subplot(2, 2, i + 1)
22
    ax.imshow(np.sqrt(np.abs(wp2[p2].data)), origin='upper',
23
              interpolation="nearest", cmap=plt.cm.gray)
24
    ax.set_title(p2)
25

26
# Show level 2 nodes
27
for p1 in path:
28
    fig = plt.figure()
29
    for i, p2 in enumerate(path):
30
        ax = fig.add_subplot(2, 2, i + 1)
31
        p1p2 = p1 + p2
32
        ax.imshow(np.sqrt(np.abs(wp2[p1p2].data)), origin='upper',
33
                  interpolation="nearest", cmap=plt.cm.gray)
34
        ax.set_title(p1p2)
35

36
fig = plt.figure()
37
i = 1
38
for row in wp2.get_level(2, 'freq'):
39
    for node in row:
40
        ax = fig.add_subplot(len(row), len(row), i)
41
        ax.set_title("%s=(%s row, %s col)" % (
42
                     (node.path,) + wp2.expand_2d_path(node.path)))
43
        ax.imshow(np.sqrt(np.abs(node.data)), origin='upper',
44
                  interpolation="nearest", cmap=plt.cm.gray)
45
        i += 1
46

47
plt.show()
48