def test_yule_walker_AR():
arsig,_,_ = ut.ar_generator(N=512)
avg_pwr = (arsig*arsig.conjugate()).mean()
w, psd = tsa.yule_AR_est(arsig, 8, 1024)
# for efficiency, let's leave out the 2PI in the numerator and denominator
# for the following integral
dw = 1./1024
avg_pwr_est = np.trapz(psd, dx=dw)
npt.assert_almost_equal(avg_pwr, avg_pwr_est, decimal=0)
import numpy as np
from matplotlib import pyplot as plt
from nitime import utils
from nitime import algorithms as alg
reload(utils)
npts = 2048*10
sigma = 1
drop_transients = 1024
coefs = np.array([0.9, -0.5])
# Generate AR(2) time series
X, v, _ = utils.ar_generator(npts, sigma, coefs, drop_transients)
# Visualize
plt.figure()
plt.plot(v)
plt.title('noise')
plt.figure()
plt.plot(X)
plt.title('AR signal')
# Estimate the model parameters
sigma_est, coefs_est = alg.yule_AR_est(X, 2, 2*npts, system=True)
print 'coefs :', coefs
print 'coefs est:', coefs_est
plt.show()
