Model Data with Multiple AR models?

Questions
1

Hi,

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[0].shape[0] = 5000, input[1].shape[0] = 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)