pywt

1
from itertools import product
2

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

6
from pywt._doc_utils import (
7
    draw_2d_fswavedecn_basis,
8
    draw_2d_wp_basis,
9
    wavedec2_keys,
10
    wavedec_keys,
11
)
12

13
shape = (512, 512)
14

15
max_lev = 4       # how many levels of decomposition to draw
16
label_levels = 2  # how many levels to explicitly label on the plots
17

18
if False:
19
    fig, axes = plt.subplots(1, 4, figsize=[16, 4])
20
    axes = axes.ravel()
21
else:
22
    fig, axes = plt.subplots(2, 2, figsize=[8, 8])
23
    axes = axes.ravel()
24

25
# plot a 5-level standard DWT basis
26
draw_2d_wp_basis(shape, wavedec2_keys(max_lev), ax=axes[0],
27
                 label_levels=label_levels)
28
axes[0].set_title(f'wavedec2 ({max_lev} level)')
29

30
# plot for the fully separable case
31
draw_2d_fswavedecn_basis(shape, max_lev, ax=axes[1], label_levels=label_levels)
32
axes[1].set_title(f'fswavedecn ({max_lev} level)')
33

34
# get all keys corresponding to a full wavelet packet decomposition
35
wp_keys = list(product(['a', 'd', 'h', 'v'], repeat=max_lev))
36
draw_2d_wp_basis(shape, wp_keys, ax=axes[2])
37
axes[2].set_title(f'wavelet packet\n(full: {max_lev} level)')
38

39
# plot an example of a custom wavelet packet basis
40
keys = ['aaaa', 'aaad', 'aaah', 'aaav', 'aad', 'aah', 'aava', 'aavd',
41
        'aavh', 'aavv', 'ad', 'ah', 'ava', 'avd', 'avh', 'avv', 'd', 'h',
42
        'vaa', 'vad', 'vah', 'vav', 'vd', 'vh', 'vv']
43
draw_2d_wp_basis(shape, keys, ax=axes[3], label_levels=label_levels)
44
axes[3].set_title('wavelet packet\n(custom)')
45

46
plt.tight_layout()
47
plt.show()
48