Hi, I am trying to fit a Poisson regression model with MCMC sampling.
My problem is that the sampling breaks down for high values of the rate of the Poisson process. Take the minimal reproducer below: if I set the mean value of my (dummy) data to 400, everything is fine and I can fit the model (top plots below). However, if I increase the value to e.g. 1000, the sampling breaks down completely, see bottom plots below. Any idea what I am doing wrong?
import numpyro as nmp
from numpyro.infer import MCMC, NUTS, Predictive
import numpyro.distributions as dist
from jax import random
import jax.numpy as jnp
from numpy.random import default_rng
import arviz as az
# Create dummy data
rng = default_rng()
mean = 400
samples = rng.poisson(mean, size=1000)
# Define model
def model(data):
mu = nmp.sample('mu', dist.Normal(8., 2.))
lam = jnp.exp(mu)
y = nmp.sample("y", dist.Poisson(lam),
obs=data)
# Fit model
nuts_kernel = NUTS(model, target_accept_prob=0.99)
mcmc = MCMC(nuts_kernel, num_samples=4000, num_warmup=2000, num_chains=1)
rng_key = random.PRNGKey(0)
mcmc.run(rng_key, samples)
az.plot_trace(mcmc)
