Create a gist now

Instantly share code, notes, and snippets.

import pylab
import numpy
PI = numpy.pi
def covariance(x, k):
N = len(x) - k
return (x[:-k] * x[k:]).sum() / N
def phd1(x):
"""
Estimate frequency using Pisarenko Harmonic Decomposition.
It returns frequency `omega` in the unit radian/steps.
If `x[n] = cos(omega*n+phi)` then it returns an estimat of `omega`.
Note that mean of `x` must be 0.
See equation (6) from [Kenneth W. K. Lui and H. C. So]_.
.. [Kenneth W. K. Lui and H. C. So] An Unbiased Pisarenko Harmonic
Decomposition Estimator For Single-Tone Frequency,
http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.75.4859
"""
r1 = covariance(x, 1)
r2 = covariance(x, 2)
a = (r2 + numpy.sqrt(r2 ** 2 + 8 * r1 ** 2)) / 4 / r1
if a > 1: # error should be raised?
a = 1
elif a < -1:
a = -1
return numpy.arccos(a)
def freq(x, sample_step=1, dt=1.0):
"""Estimate frequency using `phd1`"""
omega = phd1(x[::sample_step])
return omega / 2.0 / PI / sample_step / dt
def plot_x_and_psd_with_estimated_omega(x, sample_step=1, dt=1.0):
y = x[::sample_step]
F = freq(x, sample_step, dt)
T = 1.0 / F
pylab.clf()
# plot PSD
pylab.subplot(211)
pylab.psd(y, Fs=1.0 / sample_step / dt)
ylim = pylab.ylim()
pylab.vlines(F, *ylim)
pylab.ylim(ylim)
# plot time series
pylab.subplot(223)
pylab.plot(x)
# plot time series (three cycles)
n = int(T / dt) * 3
m = n // sample_step
pylab.subplot(224)
pylab.plot(x[:n])
pylab.plot(numpy.arange(0, n, sample_step)[:m], y[:m], 'o')
pylab.suptitle('F = %s' % F)
if __name__ == '__main__':
F = 0.01
num = int(1.0 / F) * 50
n = numpy.arange(num)
x = numpy.sin(2 * PI * F * n)
x += numpy.random.randn(len(x)) * 0.1
# NB: here, theoretically x.mean() == 0
sample_step = 20
plot_x_and_psd_with_estimated_omega(x, sample_step)
pylab.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment