Checking MCMC samples with ArviZ¶

Sampling the posterior with the No-U-Turn Sampler (NUTS) ends with a distribution of threshold contours, samples we drew from the posterior. Here we show you how to check whether the samples represent the posterior, and which sampler argument to change when they do not.

ArviZ is an established library for diagnosing Bayesian models. covariance_to_arviz hands a psyphy fit to it in the shape it expects, so every ArviZ function works on our samples with no reshaping. ArviZ and the code call these draws, as in the draw dimension and n_draws. We say samples. The figure below carries two of the core diagnostics and is a teaser of what this tutorial builds towards.

The summary check for posterior samples with ArviZ: Bulk effective sample size (ESS), tail ESS and R-hat, one panel each, over every quantity at every reference point. Dashed lines are ArviZ's reference thresholds, in total ESS across all chains. We want both ESS curves entirely right of their line and the R-hat curve entirely left of its.

Everything here runs on the fit from the sampling tutorial, from the same script, nuts_inference_example.py. You can rebuild every diagnostic and figure from the cached samples without a GPU (a laptop CPU is fine), like this:

python nuts_inference_example.py --mode full --from-fit

What to check¶

4 quantities to check. The R-hat and ESS thresholds are ArviZ's defaults, which follow Vehtari et al. (2021). The divergence and tree-depth guidance comes from Stan, a probabilistic programming language for Bayesian models whose MCMC advice is the usual reference, in its reference manual: a divergence is a validity problem (e.g. step size too large), tree-depth saturation only an efficiency one (e.g., whether we need to increase tree-depth).

check	what it detects	suspect when
divergences, `sampler_stats["diverging"]`	the integrator failed, usually in a region of high curvature	any divergence at all
R-hat	chains have not mixed to the same distribution	\(\hat{R} > 1.01\)
Effective Sample Size (ESS)	autocorrelation: how many independent samples the chain is worth	bulk or tail ESS \(< 100\) per chain, so \(< 400\) in total across 4 chains
`sampler_stats["n_steps"]`	trajectories truncated at the tree-depth cap	at the cap for most samples

BlackJAX flags a divergence on every sample it takes. NUTSSampler keeps those flags in sampler_stats, a plain dict of JAX arrays:

Divergences and tree-depth saturation
_div = nuts_posterior.sampler_stats["diverging"]
_n_div = int(jnp.sum(_div))
print(
    f"  Divergences: {_n_div} / {_div.size} samples ({_n_div / _div.size:.1%})"
    f"{'  <-- investigate; try target_acceptance_rate=0.9' if _n_div else ''}"
)
# Fraction of post-warmup samples whose trajectory was cut off by the depth cap.
# A large fraction means max_num_doublings is limiting mixing.
_n_steps = nuts_posterior.sampler_stats["n_steps"]
_at_cap = float(jnp.mean(_n_steps >= 2**MAX_NUM_DOUBLINGS - 1))
print(
    f"  Leapfrog steps/sample: mean {float(jnp.mean(_n_steps)):.0f}, "
    f"{_at_cap:.0%} at the cap ({2**MAX_NUM_DOUBLINGS - 1})"
)

A divergent transition means the integrator, the thing approximating the trajectory in small steps (sampling tutorial), left the posterior's geometry. More sampling doesn't fix the samples it produced.

Any divergence makes the run suspect

One is enough, regardless of what R-hat says, which is why we check divergences first.

The second print is the fraction of samples whose trajectory hit max_num_doublings: a large fraction means the cap is limiting mixing.

Diagnostics on \(\Sigma\), with `covariance_to_arviz`¶

Diagnostics run on \(\Sigma\), not on the raw parameters, for the reason we discussed in the sampling tutorial. One call hands the fit to ArviZ:

R-hat and ESS on the covariance field
dt_sigma = covariance_to_arviz(nuts_posterior, ref_points)
summary = az.summary(dt_sigma, var_names=["Sigma"])
print(summary[["mean", "sd", "r_hat", "ess_bulk"]].describe())
r_hat_max = float(summary["r_hat"].max())
ess_min = float(summary["ess_bulk"].min())
print(f"  R-hat max on Σ: {r_hat_max:.3f}  (target < 1.01)")
print(f"  ESS min on Σ:   {ess_min:.0f}    (target > 100 per chain)")

az.summary gives one row per quantity: two log variances and one correlation at each of the 25 reference points we chose, (75 rows total). We see the print output from above for the --mode full run below:

            mean         sd      r_hat     ess_bulk
count  75.000000  75.000000  75.000000    75.000000
mean   -2.796666   0.148029   1.004584  1012.223390
std     1.892634   0.053201   0.003398   314.699325
min    -6.242783   0.077725   0.999981   398.680579
25%    -4.171345   0.108828   1.002054   806.839604
50%    -3.697674   0.142442   1.003868  1034.450751
75%    -0.410639   0.163172   1.006310  1207.652176
max     0.333164   0.337189   1.017289  1837.560237
  R-hat max on Σ: 1.017  (target < 1.01)
  ESS min on Σ:   399    (target > 100 per chain)

If you'd like to see all 75 rows instead, just drop the .describe():

Every quantity, one row each
print(summary[["mean", "sd", "r_hat", "ess_bulk"]])

covariance_to_arviz takes the posterior and the points to evaluate \(\Sigma\) at, here the reference points. From there az.summary, az.rhat, az.ess and every ArviZ plotting function work on a psyphy fit directly, which is where the two figures below come from.

Hand it the posterior, not samples from it

posterior.sample(n) flattens (n_chains, n_draws, ...) into one axis, so a sample no longer carries the chain it came from. R-hat compares the spread between chains against the spread within them, so once the chains are gone it comes back near 1.00 however badly they disagreed. ESS is inflated the same way, since shuffling removes the autocorrelation it measures.

Interpreting the figures¶

Both figures below are one ArviZ call each, az.plot_trace and az.plot_convergence_dist, on the object covariance_to_arviz returned. (We write no additional plotting code here.)

covariance_to_arviz summarizes each \(2\times2\) \(\Sigma\) by the three numbers that determine it: log_var_0 and log_var_1, the log variances along each axis, and corr_0_1, their correlation. A trace plot shows how each of those moved over the run. If the chains converged, they should overlap, stay level, and vary fast from sample to sample.

Trace plot at the worst reference point
# pick the reference point with the worst R-hat
worst_point = int(az.rhat(dt_sigma)["Sigma"].max("quantity").argmax(dim="point"))
pc = az.plot_trace(dt_sigma, var_names=["Sigma"], coords={"point": worst_point})
pc.savefig(plot_path("nuts_trace_sigma.png"), bbox_inches="tight")

✓ Chains mixed. One panel per quantity at the reference point with the worst R-hat, one colored line per chain. Horizontal axis is the sample index, which ArviZ labels `Draw`. Vertical axis is the value of that quantity.

For comparison, below we show the trace plot of the same reference point from a run that did not converge:

✗ Chains not mixed. The same plot from `--mode quick`: 2 chains, 150 samples.

Read each panel along the horizontal axis:

in the trace	interpretation	where to see it
chains overlap, each level, values varying fast from sample to sample	✓ mixed	the `--mode full` figure
one chain holds its own level	✗ not mixed	`corr_0_1` in `--mode quick`, the two chains apart for stretches of 40 samples or more
a slope at either end	✗ still drifting, so not mixed
slow swings instead of fast variation	✗ autocorrelated, so few independent samples	all three `--mode quick` panels

A trace plot catches those. It does not catch chains that look fine and still carry too few independent samples, which is what ESS and R-hat are for.

We see one reference point here. The figure at the top covers all 25 at once.

ESS and R-hat over all quantities
pc = az.plot_convergence_dist(dt_sigma, var_names=["Sigma"])
pc.savefig(plot_path("nuts_convergence_sigma.png"), bbox_inches="tight")

The figure from the top of this page, from the code above. The summary check for posterior samples with ArviZ: Bulk effective sample size (ESS), tail ESS and R-hat, one panel each, over every quantity at every reference point. Dashed lines are ArviZ's reference thresholds, in total ESS across all chains. We want both ESS curves entirely right of their line and the R-hat curve entirely left of its.

Each curve there is the distribution of one diagnostic over every \(\Sigma\) quantity at every reference point, so a point on it reads as "this fraction of quantities scored at most this value". More is better for ESS and less is better for R-hat. For R-hat: where its curve crosses the dashed line, that fraction of quantities sits above 1.01. Here a small fraction does, which is the 1.017 the summary reported.

What to change when a check fails¶

A failing check is fixed by the sampler arguments, which blackjax.nuts and window_adaptation document.

argument	what it controls	psyphy default	raise it when
`num_warmup`	tuning steps per chain, discarded	500	trace plots still drift at the start
`num_samples`	kept samples per chain	1000	ESS is too low
`num_chains`	independent chains, R-hat needs at least 2	4	a stronger convergence check is wanted
`max_num_doublings`	tree-depth cap, at most \(2^k - 1\) gradients per sample	8, against BlackJAX's 10	most samples hit the cap, which the script prints
`target_acceptance_rate`	how cautious the tuned step size is	0.8	divergences appear

Compute presets

Compute settings
MODES = {
    "quick": {
        "mc_samples": 50,
        "trials_per_ref": 50,
        "n_ref_grid": 3,
        "map_steps": 300,
        "warmup": 100,
        "draws": 150,
        "chains": 2,
        "ensemble": 30,
    },
    "balanced": {
        "mc_samples": 500,
        "trials_per_ref": 100,
        "n_ref_grid": 5,
        "map_steps": 800,
        "warmup": 500,
        "draws": 500,
        "chains": 4,
        "ensemble": 100,
    },
    "full": {
        "mc_samples": 2000,
        "trials_per_ref": 200,
        "n_ref_grid": 5,
        "map_steps": 1500,
        "warmup": 1000,
        "draws": 1000,
        "chains": 4,
        "ensemble": 200,
    },
}

The presets scale three separate things: how much data there is, how accurately the likelihood is estimated (mc_samples), and how long the sampler runs. The first two change the posterior being targeted, so a quick fit is a smaller problem rather than a rough version of a full one.

Compute costs of a run

The Monte Carlo likelihood dominates. Each gradient evaluation simulates mc_samples noisy internal representations for every trial, so cost scales with mc_samples * n_trials. quick to full raises that product by about 440 times.
Cost per sample varies. NUTS needs \(O(2^{\text{tree depth}})\) gradient evaluations and the depth is data dependent, commonly 16 to 64.
Tree depth is capped. The default max_num_doublings=8 allows at most \(2^8 - 1\) gradient evaluations per sample. If sampler_stats["n_steps"] sits at the cap for most samples, trajectories are being truncated: still valid, but mixing more slowly.
Warmup starts from a prior-scaled metric. By default the sampler works in \(z = W / \sigma_\text{prior}\) (prior_preconditioning=True), so adaptation does not begin from an identity mass matrix mismatched to the coefficient scales.
Chains are vmapped. All chains advance in one compiled call per phase. Under vmap each step costs as much as the deepest tree across chains (which is another reason for the depth cap). On a GPU four chains cost about as much as one. On a CPU, many chains can be slower than running them one at a time.

What step_size does

step_size picks between two code paths. The paths are BlackJAX's, but the choice is psyphy's: BlackJAX has no notion of modes, and NUTSSampler dispatches on whether this argument is None.

`step_size`	what happens	call
`None` (default)	`window_adaptation` tunes the step size and mass matrix per chain, starting from a prior-scaled metric	`NUTSSampler(num_warmup=200, num_samples=500)`
a float	no adaptation and an identity mass matrix. Warmup becomes plain burn-in, so it is faster, and the step size has to be chosen by hand	`NUTSSampler(..., step_size=0.01)`

psyphy deviates in two places from the BlackJax defaults: prior_preconditioning=True samples \(z = W / \sigma_\text{prior}\) and maps the samples back to \(W\), so adaptation starts from a metric matched to the prior scales. And max_num_doublings defaults to 8 here (against BlackJAX's own 10).