1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3
4# Plot scaling and wavelet functions for db, sym, coif, bior and rbio families
5
6import itertools
7
8import matplotlib.pyplot as plt
9
10import pywt
11
12
13plot_data = [('db', (4, 3)),
14             ('sym', (4, 3)),
15             ('coif', (3, 2))]
16
17
18for family, (rows, cols) in plot_data:
19    fig = plt.figure()
20    fig.subplots_adjust(hspace=0.2, wspace=0.2, bottom=.02, left=.06,
21                        right=.97, top=.94)
22    colors = itertools.cycle('bgrcmyk')
23
24    wnames = pywt.wavelist(family)
25    i = iter(wnames)
26    for col in range(cols):
27        for row in range(rows):
28            try:
29                wavelet = pywt.Wavelet(next(i))
30            except StopIteration:
31                break
32            phi, psi, x = wavelet.wavefun(level=5)
33
34            color = next(colors)
35            ax = fig.add_subplot(rows, 2 * cols, 1 + 2 * (col + row * cols))
36            ax.set_title(wavelet.name + " phi")
37            ax.plot(x, phi, color)
38            ax.set_xlim(min(x), max(x))
39
40            ax = fig.add_subplot(rows, 2*cols, 1 + 2*(col + row*cols) + 1)
41            ax.set_title(wavelet.name + " psi")
42            ax.plot(x, psi, color)
43            ax.set_xlim(min(x), max(x))
44
45for family, (rows, cols) in [('bior', (4, 3)), ('rbio', (4, 3))]:
46    fig = plt.figure()
47    fig.subplots_adjust(hspace=0.5, wspace=0.2, bottom=.02, left=.06,
48                        right=.97, top=.94)
49
50    colors = itertools.cycle('bgrcmyk')
51    wnames = pywt.wavelist(family)
52    i = iter(wnames)
53    for col in range(cols):
54        for row in range(rows):
55            try:
56                wavelet = pywt.Wavelet(next(i))
57            except StopIteration:
58                break
59            phi, psi, phi_r, psi_r, x = wavelet.wavefun(level=5)
60            row *= 2
61
62            color = next(colors)
63            ax = fig.add_subplot(2*rows, 2*cols, 1 + 2*(col + row*cols))
64            ax.set_title(wavelet.name + " phi")
65            ax.plot(x, phi, color)
66            ax.set_xlim(min(x), max(x))
67
68            ax = fig.add_subplot(2*rows, 2*cols, 2*(1 + col + row*cols))
69            ax.set_title(wavelet.name + " psi")
70            ax.plot(x, psi, color)
71            ax.set_xlim(min(x), max(x))
72
73            row += 1
74            ax = fig.add_subplot(2*rows, 2*cols, 1 + 2*(col + row*cols))
75            ax.set_title(wavelet.name + " phi_r")
76            ax.plot(x, phi_r, color)
77            ax.set_xlim(min(x), max(x))
78
79            ax = fig.add_subplot(2*rows, 2*cols, 1 + 2*(col + row*cols) + 1)
80            ax.set_title(wavelet.name + " psi_r")
81            ax.plot(x, psi_r, color)
82            ax.set_xlim(min(x), max(x))
83
84plt.show()
85