Implementing a manual AR1 model (original) (raw)
August 30, 2023, 12:04am 1
Hi,
I am new to pymc and I am trying to understand how to build time series models. I am starting with an AR1 model where I am trying to replicate what the already available AR distribution in pymc gives us. This is to help me understand how to use scan to iterate through a time series and how I am supposed to structure the building blocks.
I had a look online at some examples and at the source code for the AR distribution but I can’t make it work. I am constantly fighting with the dimensions where the observed variable does not have the same shape as the output of my scan function. I would appreciate if someone could fix my code below or point me in the right direction.
Below there are 3 models:
- The built-in AR Model
- An ARMA model I found in the examples
- My manual AR1 model which is not working
import os
import sys
import numpy as np
import arviz as az
import pymc as pm
import pytensor
import pytensor.tensor as pt
import matplotlib.pyplot as plt
%matplotlib inline
RANDOM_SEED = 123456
np.random.seed(RANDOM_SEED)
rng = np.random.default_rng(RANDOM_SEED)
sample_size = 100 # - small so it samples fast
def simulate_ar(intercept, coef1, coef2, noise=0.3, *, warmup=10, steps=200):
# We sample some extra warmup steps, to let the AR process stabilize
draws = np.zeros(warmup + steps)
# Initialize first draws at intercept
draws[:2] = intercept
for step in range(2, warmup + steps):
draws[step] = (
intercept
+ coef1 * draws[step - 1]
+ coef2 * draws[step - 2]
+ np.random.normal(0, noise)
)
# Discard the warmup draws
return draws[warmup:]
# True parameters of the AR process
ar1_data = simulate_ar(10, -0.9, 0, steps = 200 * 2)
fig, ax = plt.subplots(figsize=(10, 3))
ax.set_title("Generated Autoregressive Timeseries", fontsize=15)
ax.plot(ar1_data)
# Using links as reference
# - https://www.pymc.io/projects/examples/en/latest/time_series/AR.html
# - https://www.pymc.io/projects/examples/en/latest/time_series/Forecasting_with_structural_timeseries.html
with pm.Model() as ar1:
y = pm.MutableData("y", ar1_data, dims="obs_id")
rho = pm.Normal("rho", mu=0.0, sigma=1.0, shape=2)
sigma = pm.HalfNormal("sigma", 5)
init = pm.Normal.dist(0, 10)
ar = pm.AR(
"ar",
rho=rho,
sigma = sigma,
constant=True,
init_dist = init,
steps = y.shape[0] - rho.shape[0] + 1
)
outcome = pm.Normal("likelihood", mu=ar, sigma=sigma, observed=y, dims="obs_id")
idata_ar = pm.sample_prior_predictive()
idata_ar = pm.sample(sample_size, tune=2000, random_seed=RANDOM_SEED)
idata_ar.extend(pm.sample_posterior_predictive(idata_ar))
# Using link as reference
# https://github.com/pymc-devs/pymc-examples/blob/main/examples/time_series/arma_example.py
with pm.Model() as arma_model:
sigma = pm.HalfNormal("sigma", 5.0)
theta = pm.Normal("theta", 0.0, sigma=1.0)
phi = pm.Normal("phi", 0.0, sigma=2.0)
mu = pm.Normal("mu", 0.0, sigma=10.0)
err0 = y[0] - (mu + phi * mu)
def calc_next(last_y, this_y, err, mu, phi, theta):
nu_t = mu + phi * last_y + theta * err
return this_y - nu_t
err, _ = pytensor.scan(
fn=calc_next,
sequences=dict(input=y, taps=[-1, 0]),
outputs_info=[err0],
non_sequences=[mu, phi, theta],
)
pm.Potential("like", pm.logp(pm.Normal.dist(0, sigma=sigma), err))
example = pm.sample(sample_size, tune = 100, random_seed = RANDOM_SEED)
# --- this one does not work
# --- trying to replicate the original implementation without all the complexity
# https://github.com/pymc-devs/pymc/blob/main/pymc/distributions/timeseries.py
with pm.Model() as ar_manual:
# assumes 95% of prob mass is between -2 and 2
rho = pm.Normal("rho", mu=0.0, sigma=1.0, shape=2)
# precision of the innovation term
sigma = pm.HalfNormal("sigma", 5)
y = pm.MutableData("y", ar1_data, dims="obs_id")
init = pm.Normal.dist(0, 10, size = 1)
def calc_next(last_y, rho):
nu_t = rho[0] + rho[1] * last_y
return nu_t
upd, _ = pytensor.scan(
fn=calc_next,
outputs_info=dict(initial = init.type(), taps = [-1]),
non_sequences=rho,
n_steps = y.shape[0] - rho.shape[0] + 1,
strict = True
)
# --- Need to concat the initial value but the shapes don't match
# upd = pt.concatenate([init, upd], axis = -1)
ar = pm.Normal("ar", mu = upd, sigma = sigma)
outcome = pm.Normal("likelihood", mu=ar, sigma=sigma, observed=y, dims="obs_id")
idata_ar = pm.sample_prior_predictive()
idata_ar = pm.sample(sample_size, tune=2000, random_seed=RANDOM_SEED)
idata_ar.extend(pm.sample_posterior_predictive(idata_ar))
Thanks
In the newer versions of pymc you can define such models generatively with CustomDist:
A few specific comments about your 3rd implementation:
The concatenate fails because you need to add a dummy batch index to init. pt.concatenate([init[None], upd], axis=0]) should work.
What you’re implementing is quite close to the OLS formulation of AR, which would look something like this:
with pm.Model(coords={'params':['intercept', 'L1.data']}) as ar1:
y = pm.MutableData("y", ar1_data[:-1], dims="obs_id")
X = pm.MutableData('X', np.c_[np.ones_like(ar1_data[1:]), ar1_data[1:]], dims=[None, 'params'])
rhos = pm.Normal('rhos', dims=['params'])
mu = pm.Deterministic('mu', X @ rhos)
sigma = pm.Exponential('sigma', 1)
y_hat = pm.Normal('y_hat', mu=mu, sigma=sigma, observed=y)
idata = pm.sample()
The problem with trying to do this via a scan is that, without any innovations inside the scan, it’s just going to converge to the steady state (assuming values of rho that induce stationarity). Here’s a demonstration:
n_lags = 1
coords = {'params': ['intercept'] + [f'L{i+1}' for i in range(n_lags)]}
with pm.Model(coords=coords) as ar1:
y = pm.MutableData("y", ar1_data, dims="obs_id")
rhos = pm.Normal('rhos', sigma=[1, 0.25], dims=['params'])
init = pm.Normal('y_init', 0, 10)
def ar_step(mu_t, rho):
mu_tp1 = rho[0] + rho[1] * mu_t
return mu_tp1
mu, _ = pytensor.scan(
ar_step,
outputs_info=[init],
non_sequences=rhos,
n_steps = y.shape[0] - n_lags,
strict = True
)
mu = pm.Deterministic('mu', pt.concatenate([init[None], mu], axis=0))
# Analytic steady-state for AR(1) with non-zero constant
mu_steady = pm.Deterministic('mu_steady', rhos[0] / (1 - rhos[1]))
sigma = pm.Exponential('sigma', 1)
y_hat = pm.Normal('y_hat', mu=mu, sigma=sigma, observed=y)
prior_idata = pm.sample_prior_predictive()
You can inspect the prior draws for mu and mu_steady to check I’m not lying to you:
prior = az.extract(prior_idata, 'prior')
idx = np.random.choice(500)
sample = prior.isel(sample=idx)
fig, ax = plt.subplots(figsize=(14,4))
ax.plot(sample.mu)
ax.axhline(sample.mu_steady, ls='--', color='k')
This is why you need to follow the code @ricardoV94 linked when you do an AR model “by hand” with scan. You need to put a prior over the entire AR trajectory, which is not the same thing as running a deterministic AR simulation, then using that as the mean of an i.i.d normal process. What would be equivalent would be to scan over the data, and essentially make a sequence of 1-step forecasts, and pass that to normal:
n_lags = 1
coords = {'params': ['intercept'] + [f'L{i+1}' for i in range(n_lags)]}
with pm.Model(coords=coords) as ar1:
y = pm.MutableData("y", ar1_data, dims="obs_id")
rhos = pm.Normal('rhos', sigma=[1, 0.25], dims=['params'])
init = pm.Normal('y_init', 0, 10)
def ar_step(mu_t, rho):
mu_tp1 = rho[0] + rho[1] * mu_t
return mu_tp1
mu, _ = pytensor.scan(
ar_step,
sequences=dict(input=y, taps=[-1]),
non_sequences=rhos,
n_steps = y.shape[0] - n_lags,
strict = True
)
mu = pm.Deterministic('mu', pt.concatenate([init[None], mu], axis=0))
mu_steady = pm.Deterministic('mu_steady', rhos[0] / (1 - rhos[1]))
sigma = pm.Exponential('sigma', 1)
y_hat = pm.Normal('y_hat', mu=mu, sigma=sigma, observed=y)
idata = pm.sample()
But this is essentially a slower version of the OLS model.
The moral of the story is that pymc-examples is desperately in need of a “how to write a recursive model” notebook. Help wanted 
FYI, the Bayesian Gods made @jessegrabowski’s wishes a reality, as I think this notebook is what he was calling for 
Broken link @AlexAndorra. I think this is the right line?
Good catch @drbenvincent !