Reparametrizing of grouped AR

Questions version agnostic
1

Hey,

I’m experimenting with a simple model to test out AR parameters and ran into the issue of divergences after moving the model from a single AR for just one class of observations in the data to one with an AR parameter per class.
I’m familiar with the centered and non-centered parametrizations for hierarchical regression models, but I’m struggling to see how I should reparametrize the below model and what exactly about it is a problem.

Can someone point out what is causing divergences in this model and how I should reparametrize it?

Code

import numpy as np
import pandas as pd
import scipy as sp
import arviz as az
import pymc as pm

rng = np.random.default_rng(0)
# generate data
x = np.arange(1,11)
mu_vector = np.array([
    np.linspace(3,7,10),
    [np.sin(_x)+5 for _x in x],
    [0.002*_x**3 + 0.003*_x**2 + 3 for _x in x]
])
y = sp.stats.norm.rvs(loc=mu_vector, scale=np.ones(mu_vector.shape)*0.5)
df = pd.DataFrame({
    'x': np.tile(x,mu_vector.shape[0]).reshape(mu_vector.shape).flatten(), 
    'y': y.flatten(),
    'c': np.array([[1]*mu_vector.shape[1],[2]*mu_vector.shape[1],[3]*mu_vector.shape[1]]).flatten()
})

# Define model variables, train & test set 
X = df['x'].values-1
y = df['y'].values
c = df['c'].values-1
train_ind = df['x']<=6
test_ind = df['x']>6

# Define model
coords={
    'season':np.arange(10),
    'class':np.unique(c),
    'ar_terms':np.arange(4)
}
with pm.Model(coords=coords) as model:
    d_season = pm.Data("d_season", X[train_ind], dims='obs_id')
    d_c = pm.Data("d_c", c[train_ind], dims='obs_id')
    coefs = pm.Normal("coefs", 0, dims=('class','ar_terms'))
    ar3 = pm.AR("ar3", coefs, sigma=0.25, init_dist=pm.Normal.dist(mu=[4,3,2],sigma=1, size=3), ar_order=3, constant=True, dims=('class','season'))
    sigma = pm.HalfNormal("sigma", sigma=1.0)
    mu = pm.Deterministic("mu",ar3[d_c, d_season], dims='obs_id')
    y_ = pm.Normal("y_", mu=mu, sigma=sigma, observed=y[train_ind], dims='obs_id')
    idata = pm.sample(3000, tune=3000, chains=4,target_accept=0.95, nuts_sampler="numpyro")
2

Ok I feel like a fool, but the issue was simply that I forgot to specify a separate sigma value for coefs and the model ended up being poorly specified with the default of 1.

2 Likes