# Generated by pandoc-plot 1.8.0
import matplotlib.pyplot as plt
import numpy as np

np.random.seed(2019)

import matplotlib.pyplot as plt
import pywt

fig, ax = plt.subplots(figsize=(8,4))

def discrete_colors(it):
    """Returns a list of discrete colors to plot, for example, various time-traces."""
    cmap = plt.get_cmap("inferno")
    mi, ma = 0.11, 0.75

    it = list(it)

    if len(it) == 1:
        yield from zip([cmap(mi)], it)
        return
    elif len(it) == 2:
        yield from zip([cmap(mi), cmap(ma)], it)
        return

    step = (ma - mi) / (len(it) - 1)
    yield from zip([cmap(mi + i * step) for i in range(len(it))], it)

offset = 0
for color, level in discrete_colors([1,2,3,4,5]):
    _, wavelet, _ = pywt.Wavelet("db4").wavefun(level=level)
    
    ax.plot(wavelet + offset, color=color, label=f"level={level}")
    offset -= 2

ax.legend()
ax.grid(True)
ax.set_ylabel('Amplitude [a.u.]')
plt.yticks(color='w') # Hiding ytick labels, which are meaningless due to the offset.
ax.set_xlabel('Time [a.u.]')