Fitting an AR model: algorithm module interface

Auto-regressive (AR) processes are processes that follow the following equation:

\[x_t = \sum_{i=1}^{n}a_i * x_{t-i} + \epsilon_t\]

In this example, we will demonstrate the estimation of the AR model coefficients and the estimation of the AR process spectrum, based on the estimation of the coefficients.

We start with imports from numpy, matplotlib and import nitime.utils as well as nitime.algorithms:

import numpy as np
from matplotlib import pyplot as plt

from nitime import utils
from nitime import algorithms as alg
from nitime.timeseries import TimeSeries
from nitime.viz import plot_tseries

We define some variables, which will be used in generating the AR process:

npts = 2048
sigma = 0.1
drop_transients = 128
Fs = 1000

In this case, we generate an order 2 AR process, with the following coefficients:

coefs = np.array([0.9, -0.5])

This generates the AR(2) time series:

X, noise, _ = utils.ar_generator(npts, sigma, coefs, drop_transients)

ts_x = TimeSeries(X, sampling_rate=Fs, time_unit='s')
ts_noise = TimeSeries(noise, sampling_rate=1000, time_unit='s')

We use the plot_tseries function in order to visualize the process:

fig01 = plot_tseries(ts_x, label='AR signal')
fig01 = plot_tseries(ts_noise, fig=fig01, label='Noise')
fig01.axes[0].legend()

Now we estimate back the model parameters, using two different estimation algorithms.

coefs_est, sigma_est = alg.AR_est_YW(X, 2)
# no rigorous purpose behind 100 transients
X_hat, _, _ = utils.ar_generator(
    N=npts, sigma=sigma_est, coefs=coefs_est, drop_transients=100, v=noise
    )
fig02 = plt.figure()
ax = fig02.add_subplot(111)
ax.plot(np.arange(100, len(X_hat) + 100), X_hat, label='estimated process')
ax.plot(X, 'g--', label='original process')
ax.legend()
err = X_hat - X[100:]
mse = np.dot(err, err) / len(X_hat)
ax.set_title('Mean Square Error: %1.3e' % mse)


plt.show()

Example source code

You can download the full source code of this example. This same script is also included in the Nitime source distribution under the doc/examples/ directory.