I’m trying to model the data with two AR1 models.
The code below is a modification of the AR1 example from the official pymc3 documentation. I create a
group_idx(with 0, 1) to separate the data into two groups. But I keep getting
ValueError: Input dimension mis-match. (input.shape = 5000, input.shape = 4999). I guess the AR model behaves differently from distributions.
What is the right way to achieve this?
import pymc3 as pm import arviz as az import numpy as np RANDOM_SEED = 8927 np.random.seed(RANDOM_SEED) 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:] group_idx = np.zeros_like(y, dtype=int) group_idx[1000:] = 1 with pm.Model() as ar1: # assumes 95% of prob mass is between -2 and 2 theta = pm.Normal("theta", 0.0, 1.0, shape=2) # precision of the innovation term tau = pm.Exponential("tau", 0.5, shape=2) # process mean center = pm.Normal("center", mu=0.0, sigma=1.0) likelihood = pm.AR1("y", k=theta[group_idx], tau_e=tau[group_idx], observed=y - center) trace = pm.sample(1000, tune=2000, init="advi+adapt_diag", random_seed=RANDOM_SEED) idata = az.from_pymc3(trace)