# AR model using Yule-Walker method

Author: Matti Pastell
Date: 14.5.2013

from scipy import signal, linalg
import numpy as np
import matplotlib.pyplot as plt

class YW(object):
"""A class to fit AR model using Yule-Walker method"""

def __init__(self, X):
self.X = X - np.mean(X)


# Calculate autocorrelation

YW method requires that we compute the sample autocorrelation function:

$r_k = \frac{1}{(n-k)\sigma^2}\sum_{t=1}^{n-k}(X_t - \mu)(X_{t+k} - \mu)$

    def autocorr(self, lag=10):
c = np.correlate(self.X, self.X, 'full')
mid = len(c)//2
acov = c[mid:mid+lag]
acor = acov/acov[0]
return(acor)


# Fit

Form the Yule-Walker equations $$r = R \Phi$$ based on sample autocorrelation $$r_k$$. Notice that the matrix R is a Toeplizt matrix and it is thus easy to form using toeplitz function from scipy.linalg.

$\begin{pmatrix} r_1\\ r_2\\ \vdots\\ r_p \end{pmatrix} = \begin{pmatrix} r_0 & r_1 & \ldots & r_{p-1} \\ r_1 & r_0 & \ldots & r_{p-2} \\ \vdots & \vdots & \ddots & \vdots \\ r_{p-1} & r_{p-2} & \ldots & r_0 \end{pmatrix} \begin{pmatrix} \phi_1\\ \phi_2\\ \vdots\\ \phi_{p}\\ \end{pmatrix}$

And solve simply using: $\Phi = R^{-1}r$

    def fit(self, p=5):
ac = self.autocorr(p+1)
R = linalg.toeplitz(ac[:p])
r = ac[1:p+1]
self.phi = linalg.inv(R).dot(r)


# Calculate and plot the spectrum

The spectrum of an AR process is given by:

$S(f) = \frac{\sigma^2}{|1 - \sum_{k=1}^{p} \phi_k e^{-2\pi ikf}|^2}$

It can be calcuted easily using scipy.signal.freqz.

    def spectrum(self):
a = np.concatenate([np.ones(1), -self.phi])
w, h = signal.freqz(1, a)
h_db = 10*np.log10(2*(np.abs(h)/len(h)))
plt.plot(w/np.pi, h_db)
plt.xlabel(r'Normalized Frequency ($\times \pi$rad/sample)')
plt.title(r'Yule-Walker Spectral Density Estimate')


# Try it out:

x = np.sin(np.linspace(0, 20))
ar1 = YW(x)
ar1.fit()
ar1.phi

array([ 1.19379795, -0.21810471, -0.12747881, -0.06257484,
-0.12929761])

ar1.spectrum()