from pylab import log10, linspace, plot, xlabel, ylabel, legend, randn, pi
import scipy.signal
from spectrum import aryule, Periodogram, arma2psd
# Create a AR model
a = [1, -2.2137, 2.9403, -2.1697, 0.9606]
# create some data based on these AR parameters
y = scipy.signal.lfilter([1], a, randn(1, 1024))
# if we know only the data, we estimate the PSD using Periodogram
p = Periodogram(y[0], sampling=2)  # y is a list of list hence the y[0]
p.plot(label='Model ouput')

# now, let us try to estimate the original AR parameters
AR, P, k = aryule(y[0], 4)
PSD = arma2psd(AR, NFFT=512)
PSD = PSD[len(PSD):len(PSD)//2:-1]
plot(linspace(0, 1, len(PSD)), 10*log10(abs(PSD)*2./(2.*pi)),
    label='Estimate of y using Yule-Walker AR(4)')
xlabel(r'Normalized frequency (\times \pi rad/sample)')
ylabel('One-sided PSD (dB/rad/sample)')
legend()