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")