import numpy as nupy
import matplotlib.pyplot as plt

def morlet(omega_a, K, t0, t, fs):
    C = 0.798*omega_a/(fs*K)
    ret = nupy.exp(-1j*omega_a*(t-t0))-nupy.exp(-K**2)
    ret *= C*nupy.exp(-omega_a**2*(t-t0)**2/(2*K)**2)
    return ret

np = 1000
fs = 100
omega = 2*nupy.pi*2
K = 4
dt = 1/fs
t = nupy.arange(np)*dt

o0 = nupy.repeat(0.75*omega, np/10)
phi = nupy.concatenate((o0,2*o0,o0,2*o0,o0,2*o0,o0,2*o0,o0,2*o0))
a0 = nupy.repeat(0.5, np/10)
amps = nupy.concatenate((a0,2*a0,a0,2*a0,a0,2*a0,a0,2*a0,a0,2*a0))
sig = amps*nupy.sin(nupy.cumsum(phi)*dt)

no = 50
omega_a = nupy.arange(50)/49.*20. + 5
amps = nupy.zeros((no, np))
for f in range(no):
    for i in range(np):
        wvt = morlet(omega_a[f], K, t[i], t, fs)
        amps[f,i] = nupy.absolute(nupy.sum(sig*nupy.conjugate(wvt)))

plt.subplot(2,1,1)
plt.plot(t, sig)
plt.xlim([0,10])
plt.subplots_adjust(bottom=0.1, right=0.8, top=0.9)

plt.subplot(2,1,2)
plt.pcolor(t, omega_a, amps, cmap=plt.cm.jet)
plt.xlim([0,10])
plt.subplots_adjust(bottom=0.1, right=0.8, top=0.9)

cax = plt.axes([0.85, 0.1, 0.075, 0.37])
plt.colorbar(cax=cax)

plt.show()