goertzel.py

goertzel.py

"""goertzel algorithm python implementation
https://en.wikipedia.org/wiki/Goertzel_algorithm
"""
import numpy as np
import matplotlib.pyplot as plt


def goertzel(xs, f, fs, N):
	df = fs / N
	k = int(f/df)
	#k = int(f/fs*N)
	omega = 2.0 * np.pi * (k/N)
	ss = np.zeros(3)
	ys = np.zeros_like(xs,dtype=np.complex128)
	for n,x in enumerate(xs):
		ss[0] = x + 2.0*np.cos(omega)*ss[1]-ss[2]
		ys[n] = ss[0] - np.exp(-1.j * omega) * ss[1]
		ss[1:] = ss[:-1]
	return ys

gen_sin = lambda sec ,f, fs: np.sin(2.0*np.pi*(f/fs)*np.arange(int(sec*fs)))

fs = 6000
sec = 60.0
pt = int(fs*sec)
nfft = 256*4

xs = 0.1*np.random.randn(pt)
xs[40*fs:41*fs] += 0.2 * gen_sin(1.0,500.0,fs)
ys = goertzel(xs,500.0,fs,nfft)
ys2 = goertzel(xs,496.0,fs,nfft)
ys3 = goertzel(xs,504.0,fs,nfft)

fig = plt.figure(1)
ax = fig.add_subplot(211)
ax.plot(xs)
ax.set_ylim((-1.0,1.0))
ax = fig.add_subplot(212)
ax.plot(np.abs(ys))
ax.plot(np.abs(ys2))
ax.plot(np.abs(ys3))

plt.show()