Hey,
I’m looking to try and add an autoregressive noise term to a hierarchical / multilevel model, and do so across multiple levels in the forecast.
For example say I’ve got data at the level of:
- Region
- Item
And I’d like to have coefficients AR distributed noise for each combination of those categories. I started with something like this (borrowing from this thread on fitting a MA(1) time series), but unfortunately it’s extremely slow, I was wondering if there was a better / more effective way of doing this?
T = y.shape[0]
## Residuals component
with numpyro.plate('n_regions', n_regions):
with numpyro.plate('n_items', n_items):
rho = numpyro.sample("rho", dist.Uniform(-1., 1.))
ar_sigma = numpyro.sample("ar_sigma", dist.HalfNormal(.1))
def transition_fn(carry, t):
y_prev = carry
eps_t = rho * y_prev
eps = numpyro.sample("eps", dist.Normal(eps_t, ar_sigma))
return eps_t, eps
init_carry = jnp.empty(shape=(n_items, n_regions), dtype=jnp.float32)
with numpyro.handlers.condition(data={"y": y[1:]}):
_, epsilon = scan(
transition_fn,
init_carry,
jnp.arange(0, T)
)
## Forecast
mu += epsilon
And while this individual component will run (as it’s own model), when I try and add it on to the rest of the time series which is shape (5005,) I encounter: ValueError: Incompatible shapes for broadcasting: shapes=[(2,), (6,)].
Does anyone know a way around this?