Visual MCMC diagnostics using the bayesplot package

mcmc_parcoord

The mcmc_parcoord plot shows one line per iteration, connecting the parameter values at this iteration. This lets you see global patterns in the divergences.

This function works in general without including information about the divergences, but if the optional np argument is used to pass NUTS parameter information, then divergences will be colored in the plot (by default in red).

color_scheme_set("darkgray")
mcmc_parcoord(posterior_cp, np = np_cp)

Here, you may notice that divergences in the centered parameterization happen exclusively when tau, the hierarchical standard deviation, goes near zero and the values of the thetas are essentially fixed. This makes tau immediately suspect. See Gabry et al. (2019) for another example of the parallel coordinates plot.

mcmc_pairs

The mcmc_pairs function can also be used to look at multiple parameters at once, but unlike mcmc_parcoord (which works well even when including several dozen parameters) mcmc_pairs is more useful for up to ~8 parameters. It shows univariate histograms and bivariate scatter plots for selected parameters and is especially useful in identifying collinearity between variables (which manifests as narrow bivariate plots) as well as the presence of multiplicative non-identifiabilities (banana-like shapes).

Let’s look at how tau interacts with other variables, using only one of the thetas to keep the plot readable:

mcmc_pairs(posterior_cp, np = np_cp, pars = c("mu","tau","theta[1]"),
           off_diag_args = list(size = 0.75))

Note that each bivariate plot is present twice – by default each of those contain half of the chains, so you also get to see if the chains produced similar results (see the documentation for the condition argument for other options). Here, the interaction of tau and theta[1] seems most interesting, as it concentrates the divergences into a tight region.

Further examples of pairs plots and instructions for using the various optional arguments to mcmc_pairs are provided via help("mcmc_pairs").

mcmc_scatter

Using the mcmc_scatter function (with optional argument np) we can look at a single bivariate plot to investigate it more closely. For hierarchical models, a good place to start is to plot a “local” parameter (theta[j]) against a “global” scale parameter on which it depends (tau).

We will also use the transformations argument to look at the log of tau, as this is what Stan is doing under the hood for parameters like tau that have a lower bound of zero. That is, even though the draws for tau returned from Stan are all positive, the parameter space that the Markov chains actual explore is unconstrained. Transforming tau is not strictly necessary for the plot (often the plot is still useful without it) but plotting in the unconstrained is often even more informative.

First the plot for the centered parameterization:

# assign to an object so we can reuse later
scatter_theta_cp <- mcmc_scatter(
  posterior_cp, 
  pars = c("theta[1]", "tau"), 
  transform = list(tau = "log"), # can abbrev. 'transformations'
  np = np_cp, 
  size = 1
)
scatter_theta_cp

The shape of this bivariate distribution resembles a funnel (or tornado). This one in particular is essentially the same as an example referred to as Neal’s funnel (details in the Stan manual) and it is a clear indication that the Markov chains are struggling to explore the tip of the funnel, which is narrower than the rest of the space.

The main problem is that large steps are required to explore the less narrow regions efficiently, but those steps become too large for navigating the narrow region. The required step size is connected to the value of tau. When tau is large it allows for large variation in theta (and requires large steps) while small tau requires small steps in theta.

The non-centered parameterization avoids this by sampling the eta parameter which, unlike theta, is a priori independent of tau. Then theta is computed deterministically from the parameters eta, mu and tau afterwards. Here’s the same plot as above, but with eta[1] from non-centered parameterization instead of theta[1] from the centered parameterization:

scatter_eta_ncp <- mcmc_scatter(
  posterior_ncp, 
  pars = c("eta[1]", "tau"), 
  transform = list(tau = "log"), 
  np = np_ncp, 
  size = 1
)
scatter_eta_ncp

We can see that the funnel/tornado shape is replaced by a somewhat Gaussian blob/cloud and the divergences go away. Gabry et al. (2019) has further discussion of this example.

Ultimately we only care about eta insofar as it enables the Markov chains to better explore the posterior, so let’s directly examine how much more exploration was possible after the reparameterization. For the non-centered parameterization we can make the same scatterplot but use the values of theta[1] = mu + eta[1] * tau instead of eta[1]. Below is a side by side comparison with the scatterplot of theta[1] vs log(tau) from the centered parameterization that we made above. We will also force the plots to have the same y-axis limits, which will make the most important difference much more apparent:

# A function we'll use several times to plot comparisons of the centered 
# parameterization (cp) and the non-centered parameterization (ncp). See
# help("bayesplot_grid") for details on the bayesplot_grid function used here.
compare_cp_ncp <- function(cp_plot, ncp_plot, ncol = 2, ...) {
  bayesplot_grid(
    cp_plot, ncp_plot, 
    grid_args = list(ncol = ncol),
    subtitles = c("Centered parameterization", 
                  "Non-centered parameterization"),
    ...
  )
}


scatter_theta_ncp <- mcmc_scatter(
  posterior_ncp, 
  pars = c("theta[1]", "tau"), 
  transform = list(tau = "log"), 
  np = np_ncp, 
  size = 1
)

compare_cp_ncp(scatter_theta_cp, scatter_theta_ncp, ylim = c(-8, 4))

Once we transform the eta values into theta values we actually see an even more pronounced funnel/tornado shape than we have with the centered parameterization. But this is precisely what we want! The non-centered parameterization allowed us to obtain draws from the funnel distribution without having to directly navigate the curvature of the funnel. With the centered parameterization the chains never could make it into the neck of funnel and we see a clustering of divergences and no draws in the tail of the distribution.

mcmc_trace

Another useful diagnostic plot is the trace plot, which is a time series plot of the Markov chains. That is, a trace plot shows the evolution of parameter vector over the iterations of one or many Markov chains. The np argument to the mcmc_trace function can be used to add a rug plot of the divergences to a trace plot of parameter draws. Typically we can see that at least one of the chains is getting stuck wherever there is a cluster of many red marks.

Here is the trace plot for the tau parameter from the centered parameterization:

color_scheme_set("mix-brightblue-gray")
mcmc_trace(posterior_cp, pars = "tau", np = np_cp) + 
  xlab("Post-warmup iteration")

The first thing to note is that all chains seem to be exploring the same region of parameter values, which is a good sign. But the plot is too crowded to help us diagnose divergences. We may however zoom in to investigate, using the window argument:

mcmc_trace(posterior_cp, pars = "tau", np = np_cp, window = c(300,500)) + 
  xlab("Post-warmup iteration")

What we see here is that chains can get stuck as tau approaches zero and spend substantial time in the same region of the parameter space. This is just another indication that there is problematic geometry at τ ≃ 0 – healthy chains jump up and down frequently.

mcmc_nuts_divergence

To understand how the divergences interact with the model globally, we can use the mcmc_nuts_divergence function:

color_scheme_set("red")
mcmc_nuts_divergence(np_cp, lp_cp)

In the top panel we see the distribution of the log-posterior when there was no divergence vs the distribution when there was a divergence. Divergences often indicate that some part of the posterior isn’t being explored and the plot confirms that lp|Divergence indeed has lighter tails than lp|No divergence.

The bottom panel shows the same thing but instead of the log-posterior the NUTS acceptance statistic is shown.

Specifying the optional chain argument will overlay the plot just for a particular Markov chain on the plot for all chains combined:

mcmc_nuts_divergence(np_cp, lp_cp, chain = 4)

For the non-centered parameterization we may get a few warnings about divergences but if we do we’ll have far fewer of them to worry about.

mcmc_nuts_divergence(np_ncp, lp_ncp)

If there are only a few divergences we can often get rid of them by increasing the target acceptance rate (adapt_delta, the upper limit is 1), which has the effect of lowering the step size used by the sampler and allowing the Markov chains to explore more complicated curvature in the target distribution.

fit_cp_2 <- sampling(schools_mod_cp, data = schools_dat,
                     control = list(adapt_delta = 0.999), seed = 978245244)
fit_ncp_2 <- sampling(schools_mod_ncp, data = schools_dat,
                      control = list(adapt_delta = 0.999), seed = 843256842)

For the first model and this particular data, increasing adapt_delta will not solve the problem and a reparameterization is required.

mcmc_nuts_divergence(nuts_params(fit_cp_2), log_posterior(fit_cp_2))
mcmc_nuts_divergence(nuts_params(fit_ncp_2), log_posterior(fit_ncp_2))

mcmc_nuts_energy

The plot created by mcmc_nuts_energy shows overlaid histograms of the (centered) marginal energy distribution π_E and the first-differenced distribution π_ΔE,

color_scheme_set("red")
mcmc_nuts_energy(np_cp)

The two histograms ideally look the same (Betancourt, 2017), which is only the case for the non-centered parameterization (right):

compare_cp_ncp(
  mcmc_nuts_energy(np_cp, binwidth = 1/2),
  mcmc_nuts_energy(np_ncp, binwidth = 1/2)
)

The difference between the parameterizations is even more apparent if we force the step size to a smaller value and help the chains explore more of the posterior:

np_cp_2 <- nuts_params(fit_cp_2)
np_ncp_2 <- nuts_params(fit_ncp_2)

compare_cp_ncp(
  mcmc_nuts_energy(np_cp_2), 
  mcmc_nuts_energy(np_ncp_2)
)

See Betancourt (2017) for more on this particular example as well as the general theory behind the energy plots.

mcmc_rhat, mcmc_rhat_hist

We can visualize the R̂ values with the mcmc_rhat function:

color_scheme_set("brightblue") # see help("color_scheme_set")
mcmc_rhat(rhats)

In the plot, the points representing the R̂ values are colored based on whether they are less than 1.05, between 1.05 and 1.1, or greater than 1.1. There is no theoretical reason to trichotomize R̂ values using these cutoffs, so keep in mind that this is just a heuristic.

The y-axis text is off by default for this plot because it’s only possible to see the labels clearly for models with very few parameters. We can see the names of the parameters with the concerning R̂ values using the yaxis_text convenience function (which passes arguments like hjust to ggplot2::element_text):

mcmc_rhat(rhats) + yaxis_text(hjust = 1)

If we look at the same model fit using longer Markov chains we should see all R̂ < 1.1, and all points in the plot the same (light) color:

mcmc_rhat(rhat = rhat(fit_cp)) + yaxis_text(hjust = 0)

We can see the same information shown by mcmc_rhat but in histogram form using the mcmc_rhat_hist function. See the Examples section in help("mcmc_rhat_hist") for examples.

mcmc_neff, mcmc_neff_hist

ratios_cp <- neff_ratio(fit_cp)
print(ratios_cp)
mcmc_neff(ratios_cp, size = 2)

In the plot, the points representing the values of n_eff/N are colored based on whether they are less than 0.1, between 0.1 and 0.5, or greater than 0.5. These particular values are arbitrary in that they have no particular theoretical meaning, but a useful heuristic is to worry about any n_eff/N less than 0.1.

One important thing to keep in mind is that these ratios will depend not only on the model being fit but also on the particular MCMC algorithm used. One reason why we have such high ratios of n_eff to N is that the No-U-Turn sampler used by rstan generally produces draws from the posterior distribution with much lower autocorrelations compared to draws obtained using other MCMC algorithms (e.g., Gibbs).

Even for models fit using rstan the parameterization can make a big difference. Here are the n_eff/N plots for fit_cp and fit_ncp side by side.

neff_cp <- neff_ratio(fit_cp, pars = c("theta", "mu", "tau"))
neff_ncp <- neff_ratio(fit_ncp, pars = c("theta", "mu", "tau"))
compare_cp_ncp(mcmc_neff(neff_cp), mcmc_neff(neff_ncp), ncol = 1)

Because of the difference in parameterization, the effective sample sizes are much better for the second model, the non-centered parameterization.

mcmc_acf, mcmc_acf_bar

Here we can again see a difference when comparing the two parameterizations of the same model. For model 1, θ₁ is the primitive parameter for school 1, whereas for the non-centered parameterization in model 2 the primitive parameter is η₁ (and θ₁ is later constructed from η₁, μ, and τ):

compare_cp_ncp(
  mcmc_acf(posterior_cp, pars = "theta[1]", lags = 10),
  mcmc_acf(posterior_ncp, pars = "eta[1]", lags = 10)
)