Multilevel AR(1) Residuals in Time Series


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(
        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?

It seems something is missing. I’m seeing you are conditioning the data y but your transition function only includes eps variable.