goertzel.py

import math
import numpy
import matplotlib.pyplot as plt

def sine(frequency, length, rate):
    length = int(length * rate)
    factor = float(frequency) * (math.pi * 2) / rate
    return numpy.sin(numpy.arange(length) * factor)


def plotter(signal, fft):
    frq, Z = fft
    f, ax = plt.subplots(2)
    ax[0].plot(signal)
    ax[1].plot(frq, abs(Z))
    plt.show()


def compute_fft(data, sample_rate):
    n = len(data)
    Fs = sample_rate
    T = n / Fs
    k = numpy.arange(n)
    frq = k / T
    frq = frq[range(n / 2)]
    # normalize magnitudes over the number of samples the x2 factor is to
    # account for that the power is divided between the positive and negative
    # sides and we're only looking at one side
    Y = 2.0 * numpy.fft.fft(data) / n
    Z = Y[range(n / 2)]
    return (frq, Z)


def goertzel(data, freq, rate):
    prev   = 0.0
    prev2 = 0.0
    n_freq = float(freq) / rate
    coeff = 2 * math.cos(2 * math.pi * n_freq)
    for i in range(len(data)):
        s = data[i] + coeff * prev - prev2
        prev2 = prev
        prev = s
    power = prev2 * prev2 + prev * prev - coeff * prev * prev2
    return math.sqrt(power)


if __name__ == '__main__':
    rate = 48000
    chunks = sine(500, 5, rate)

    gz = [2.0 * goertzel(chunks, x, rate)/float(len(chunks)) for x in range(100, 2000, 100)]
    print gz
    x, y = compute_fft(chunks, rate)
    plotter(chunks, (x,y))