Hello. I am experimenting with an MA(1) model. I would like to fit the model to multiple time series simultaneously, but I struggle to get MCMC to converge when I include too many time series.
This issue is very similar to the one which was discussed in this topic: Expanding SGT model to multiple time series, but the solution presented there does not seem to work in my case. I may not have understood the solution well, so I would like to revisit it here.
I am running this code:
import numpyro
from numpyro.infer import MCMC, NUTS
import numpyro.distributions as dist
from jax import random
from numpyro.contrib.control_flow import scan
import jax.numpy as jnp
rng_key = random.PRNGKey(12)
seeds = random.split(rng_key, 10)
n_series = 10
series_to_include = [0,1,2,3,4,5,6,7,8,9]
def model(y):
n_series = y.shape[1]
with numpyro.plate("plate", n_series):
theta = numpyro.sample("theta", dist.Normal(0.0, 0.1))
def transition_fn(carry, t_idx):
epsilon_prev = carry
mean = theta*epsilon_prev
y_pred = numpyro.sample("y_pred", dist.Normal(mean, 1.0))
epsilon_t = y_pred - mean
carry_new = epsilon_t
return carry_new, y_pred
time_indices = jnp.arange(0, y.shape[0])
init_carry = jnp.empty(shape=(n_series), dtype=float)
with numpyro.handlers.condition(data={"y_pred": y}):
_, y_predictions = scan(
transition_fn,
init_carry,
time_indices,
)
nuts_kernel = NUTS(model)
train = random.normal(shape=(200,n_series), key = seeds[0])
mcmc = MCMC(nuts_kernel, num_samples=4000, num_warmup=2000, num_chains=1)
mcmc.run(seeds[1], y=train[:,series_to_include])
mcmc.print_summary()
I receive the following output:
mean std median 5.0% 95.0% n_eff r_hat
theta[0] 0.85 0.00 0.85 0.85 0.85 0.50 1.00
theta[1] 1.00 0.00 1.00 1.00 1.00 0.50 1.00
theta[2] 0.50 0.00 0.50 0.50 0.50 0.50 1.00
theta[3] -0.32 0.00 -0.32 -0.32 -0.32 0.50 1.00
theta[4] 0.75 0.00 0.75 0.75 0.75 0.50 1.00
theta[5] 0.71 0.00 0.71 0.71 0.71 0.50 1.00
theta[6] -1.07 0.00 -1.07 -1.07 -1.07 0.50 1.00
theta[7] -1.07 0.00 -1.07 -1.07 -1.07 0.50 1.00
theta[8] -0.24 0.00 -0.24 -0.24 -0.24 0.50 1.00
theta[9] 1.10 0.00 1.10 1.10 1.10 0.50 1.00
Number of divergences: 1009
The n_eff=0.50 values indicate that something has gone wrong. Also the correct values of theta are zero. If I remove any of the 10 series from the training data then the fitting succeeds.
In Expanding SGT model to multiple time series there is a suggestion to add .to_event(1) to the sample inside the transition function. I tried this with the below, but it didn’t solve the issue.
y_pred = numpyro.sample("y_pred", dist.Normal(mean, 1.0).to_event(1))
Does anybody have any other pointers? Thank you.
numpyro==0.10.0