Analysis of An \(AR(1)\) Model in PyMC3#

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc3 as pm

%config InlineBackend.figure_format = 'retina'
RANDOM_SEED = 8927
np.random.seed(RANDOM_SEED)
az.style.use("arviz-darkgrid")

Consider the following AR(1) process, initialized in the infinite past: $\( y_t = \theta y_{t-1} + \epsilon_t, \)\( where \)\epsilon_t \overset{iid}{\sim} {\cal N}(0,1)\(. Suppose you'd like to learn about \)\theta\( from a a sample of observations \)Y^T = { y_0, y_1,\ldots, y_T }$.

First, let’s generate some synthetic sample data. We simulate the ‘infinite past’ by generating 10,000 samples from an AR(1) process and then discarding the first 5,000:

T = 10000
y = np.zeros((T,))

# true stationarity:
true_theta = 0.95
# true standard deviation of the innovation:
true_sigma = 2.0
# true process mean:
true_center = 0.0

for t in range(1, T):
    y[t] = true_theta * y[t - 1] + np.random.normal(loc=true_center, scale=true_sigma)

y = y[-5000:]
plt.plot(y, alpha=0.8)
plt.xlabel("Timestep")
plt.ylabel("$y$");
../_images/9c9c24e849ce5dbd0387c9bfaf5be7e506d4a6b991b218f60f4f849d4c326404.png

This generative process is quite straight forward to implement in PyMC3:

with pm.Model() as ar1:
    # assumes 95% of prob mass is between -2 and 2
    theta = pm.Normal("theta", 0.0, 1.0)
    # precision of the innovation term
    tau = pm.Exponential("tau", 0.5)
    # process mean
    center = pm.Normal("center", mu=0.0, sigma=1.0)

    likelihood = pm.AR1("y", k=theta, tau_e=tau, observed=y - center)

    trace = pm.sample(1000, tune=2000, init="advi+adapt_diag", random_seed=RANDOM_SEED)
    idata = az.from_pymc3(trace)
Auto-assigning NUTS sampler...
Initializing NUTS using advi+adapt_diag...
Convergence achieved at 13500
Interrupted at 13,499 [6%]: Average Loss = 42,118
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [center, tau, theta]
Sampling 2 chains for 2_000 tune and 1_000 draw iterations (4_000 + 2_000 draws total) took 16 seconds.

We can see that even though the sample data did not start at zero, the true center of zero is captured rightly inferred by the model, as you can see in the trace plot below. Likewise, the model captured the true values of the autocorrelation parameter 𝜃 and the innovation term \(\epsilon_t\) (tau in the model) – 0.95 and 1 respectively).

az.plot_trace(
    idata,
    lines=[
        ("theta", {}, true_theta),
        ("tau", {}, true_sigma**-2),
        ("center", {}, true_center),
    ],
);
../_images/f2c25c7be199bd7657f534f9940f5128585888967fcd25974b8514782809a304.png

Extension to AR(p)#

We can instead estimate an AR(2) model using PyMC3.

\[ y_t = \theta_1 y_{t-1} + \theta_2 y_{t-2} + \epsilon_t. \]

The AR distribution infers the order of the process thanks to the size the of rho argmument passed to AR.

We will also use the standard deviation of the innovations (rather than the precision) to parameterize the distribution.

with pm.Model() as ar2:
    theta = pm.Normal("theta", 0.0, 1.0, shape=2)
    sigma = pm.HalfNormal("sigma", 3)
    likelihood = pm.AR("y", theta, sigma=sigma, observed=y)

    trace = pm.sample(
        1000,
        tune=2000,
        random_seed=RANDOM_SEED,
    )
    idata = az.from_pymc3(trace)
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [sigma, theta]
Sampling 2 chains for 2_000 tune and 1_000 draw iterations (4_000 + 2_000 draws total) took 22 seconds.
az.plot_trace(
    idata,
    lines=[
        ("theta", {"theta_dim_0": 0}, true_theta),
        ("theta", {"theta_dim_0": 1}, 0.0),
        ("sigma", {}, true_sigma),
    ],
);
../_images/55f98c1a255164a35431b99033d125cd5608d8b7d35fa32766af602798611f99.png

You can also pass the set of AR parameters as a list.

with pm.Model() as ar2_bis:
    beta0 = pm.Normal("theta0", mu=0.0, sigma=1.0)
    beta1 = pm.Uniform("theta1", -1, 1)
    sigma = pm.HalfNormal("sigma", 3)
    likelhood = pm.AR("y", [beta0, beta1], sigma=sigma, observed=y)

    trace = pm.sample(
        1000,
        tune=2000,
        random_seed=RANDOM_SEED,
    )
    idata = az.from_pymc3(trace)
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [sigma, theta1, theta0]
Sampling 2 chains for 2_000 tune and 1_000 draw iterations (4_000 + 2_000 draws total) took 25 seconds.
The acceptance probability does not match the target. It is 0.9000428081776881, but should be close to 0.8. Try to increase the number of tuning steps.
az.plot_trace(
    idata,
    lines=[("theta0", {}, true_theta), ("theta1", {}, 0.0), ("sigma", {}, true_sigma)],
);
../_images/231a32728941e676dab3ea6adf62bf1237f91bfc3594c64aafa8b7069c036fd9.png
%load_ext watermark
%watermark -n -u -v -iv -w
numpy 1.18.1
arviz 0.8.3
pymc3 3.9.3
last updated: Fri Aug 21 2020 

CPython 3.8.2
IPython 7.13.0
watermark 2.0.2