"""
Spindles-like Morlet wavelet.
Useful for automatic spindles detection.
---
Author: Raphael Vallat <raphaelvallat9@gmail.com>
Date: September 14, 2018
"""
import numpy as np
import matplotlib.pyplot as plt
from mne.time_frequency import morlet

# Parameters
sf = 100 	# Sampling rate of the signal
cf = 13.5	# Central spindles frequency
nc = 12 	# Number of cycles in the spindles

# Compute the wavelet and its envelope
wlt = morlet(sf, [cf], n_cycles=nc, zero_mean=True)[0]
wlt_env = np.abs(wlt)

# Plot
t = np.arange(wlt.size) / sf
plt.plot(t, wlt)
plt.plot(t, wlt_env, 'k')
plt.xlabel('Time (seconds)')
plt.ylabel('Amplitude')
plt.show()