wavelet_magnitude.py

import mne
import numpy as np
from scipy.signal import freqz
import matplotlib.pyplot as plt
freqs = np.arange(2, 81)
n_cycles = freqs / 2.
sfreq = 1000.
fig, ax = plt.subplots(1)
lengths = list()
for wavelet in mne.time_frequency.morlet(sfreq, freqs, n_cycles,
                                         zero_mean=True):
    wavelet = wavelet.real
    lengths.append(len(wavelet))
    w, h = freqz(wavelet, worN=2 ** 12)
    w *= sfreq / np.pi / 2.
    h = 20 * np.log10(np.abs(h))
    h -= h.max()
    ax.fill_between(w, h, -200, color='k', alpha=0.1)
xticks = freqs
ax.set(ylim=[-20, 0], xlim=[0.5, 15],
       ylabel='Response (dB)', xlabel='Frequency (Hz)',
       title='Frequency responses', xticks=xticks)
fig.tight_layout()
fig.savefig('wavelets.png')
plt.show()