Hi there, I’m trying to show the warmup steps in the ArviZ traceplot, obtained using the function arviz.plots.traceplot.
By look the the function definition, available here, I’ve realized that this cannot be done as the function is hardcoded to use the samples on the posterior distribution only.
You can see that in a line of code towards the bottom of the file that says
coords_data = get_coords(convert_to_dataset(data, group="posterior"), coords)
What group would correspond to getting the posterior plus the warmup steps?
Wouldn’t it be a good idea to add a warmup optional flag to the function call to have both the warmup and the samples on the posterior distribution?
I wouldn’t mind submitting a pull request with such a feature, if deemed worthy, because it seems quite trivial to implement.
have you taken a look at the new MCMC-monitor project?
if you’re running CmdStanPy, you can do this by specifying an output directory and setting save_warmup=True
then you can point the MCMC monitor at the output directory and you’ll not only get warmup traceplots, you’ll get it as the sampler is running. how cool is that?
1 Like
No I have not, because I didn’t know such a thing existed.
Looks kickass, thanks for sharing!
However, given how simple adding this feature seems to be, I think it would be a nice addition to ArviZ.
perhaps ArviZ can do this.
what do you want to learn from the traceplots?
My hope was to have the warmup showing in the traceplot, more for pedagogical purposes than anything else really, I just think it would be nice to show a plot for both my professors and my colleagues (who aren’t statisticians) that there was in fact a part where the chain was converging, and one it has converged.
It could also be useful to have some intuition on how much warmup steps would be required to actually get a specific model and a specific dataset to converge. If you discard it, you might have taken way too many steps.
I was looking for a simple solution of having a flag show_warmup in the trace plot function.
the default number of warmup iterations is both too high and too low, depending on the model and data.
traceplots only tell part of the story. during warmup the algorithm is tuning the stepsize and metric. if you look at the traceplot of total joint log prob for a well-specified model, the chains appear to converge very quickly - within 50 or 100 iterations. after this point, the algorithm continues to refine both stepsize and metric, and this doesn’t really show up in the traceplots, but it does lead to fewer post-warmup divergences. if the pattern of post-warmup divergences is relatively benign - i.e., not concentrated in a particular part of the posterior - then the tradeoff between fewer overall iterations vs. divergences may be worth it.
I see, then I suppose that using the tool you previously mentioned is better, since it seems to include both the traceplot and more.
I will try to set it up myself and explore what it provides me with, thanks!
1 Like
Just to add another option, arviz has a group called warmup_posterior which you can use to make your own traceplot manually using matplotlib.
Here’s an example:
import numpy as np
import matplotlib.pyplot as plt
import arviz as az
import cmdstanpy
STAN_FILE = "bernoulli.stan"
N = 100
THETA = 0.1
DATA = {
"N": N,
"y": np.random.binomial(n=1, p=THETA, size=N)
}
model_str = """
data {
int<lower=0> N;
int<lower=0,upper=1> y[N];
}
parameters {
real<lower=0,upper=1> theta;
}
model {
theta ~ beta(1,1); // uniform prior on interval 0,1
y ~ bernoulli(theta);
}
"""
with open(STAN_FILE, "w") as f:
f.write(model_str)
model = cmdstanpy.CmdStanModel(stan_file=STAN_FILE)
mcmc = model.sample(data=DATA, save_warmup=True)
idata = az.from_cmdstanpy(mcmc, save_warmup=True)
f, ax = plt.subplots(figsize=[10, 5])
traceplots = ax.plot(idata.warmup_posterior["theta"].T) # NB need to transpose!
ax.set_ylim(0, 1)
hline = ax.axhline(THETA, color="black", linestyle="--")
text = ax.set(title="Warmup traceplot", xlabel="draw", ylabel="Parameter value")
leg = f.legend(
[*traceplots, hline],
[*[f"chain {i}" for i in range(4)], "True value"],
frameon=False
)
f.savefig("warmup traceplot.png", bbox_inches="tight")
Running this script produced the graph below: you can kind of see the convergence process you want to illustrate, thought it’s maybe not the best example as it all happens in the first couple of iterations!
3 Likes
