Can mcmc.get_samples() return entire chains WITH warm-up draws?

numpyro
1

Really dumb question but I’m currently doing MCMC with AIES (or NUTS):

mcmc = MCMC(AIES(model),num_warmup=500,num_samples=500,num_chains=20,chain_method='vectorized')

mcmc.run(key(1), model_params=model_params, obs_data=obs_data, obs_sigma=0.1)

Then I can use Predictive to get posterior samples of my model parameters but this seems to exclude warmup samples by default :

Predictive(model, mcmc.get_samples())(key(1), model_params)["obs"]

Is there anyway to get the entire chain INCLUDING the warm-up samples? I want to make a movie showing how initially the draws from the prior don’t match obs_data but then eventually later samples from the chain (which may happen before/after num_warmup draws…) start to match the data. But for this I need samples from the warmup/burn-in stage…

2

Hmm, now that I think about I guess I can just set num_warmup = 0 and then num_samples = 1000 for example which means the first 500 draws would not get discarded right? I guess this should work BUT it’s not ideal since then any convergence statistics like Rhat, effective sample size, etc. computed by numpyro might be suspect. I guess in this case the user would have to do any discard of early burn-in draws themselves before computing convergence, plotting posteriors, etc. Am I getting this right? I guess implementing something like discard_warmup = True/False for mcmc.run like pymc could be an easy and useful convenience feature?

3

you can use mcmc.wamup(..., collect_warmup=True) then get samples.

1 Like