Diagnose draws with ArviZ

ArviZ summarizes and diagnoses posterior draws held in a DataTree with named variables and chain and draw dimensions. names and shapes from flatppl.Module.logdensity() name the variables, and constrain turns sampler coordinates into model parameters. Every sampler page ends here; this one samples with BlackJAX NUTS and checks ArviZ’s summary against the closed-form posterior of a Normal-Gamma model.

Runtime budget: 120 s.

from math import lgamma, log, pi
from typing import Any

import arviz as az
import blackjax
import jax
import jax.numpy as jnp
import numpy as np

from flatppl import flatppl

The model and its closed-form posterior

tau | y ~ Gamma(a_n, b_n); mu has a Student-t marginal with mean m_n and variance b_n / (k_n (a_n - 1)).

data = np.array([1.2, 0.4, 2.1, 1.7, 0.9, 1.5, 0.3, 1.1])
model = flatppl(
    """
    tau ~ Gamma(shape = 3.0, rate = 2.0)
    mu ~ Normal(mu = 0.0, sigma = 1.0 / sqrt(2.0 * tau))
    y ~ iid(Normal(mu = mu, sigma = 1.0 / sqrt(tau)), 8)
    observed = external(cartpow(reals, 8))
    prior = lawof(record(mu = mu, tau = tau))
    L = likelihoodof(kernelof(record(y = y), mu = mu, tau = tau), record(y = observed))
    posterior = bayesupdate(L, prior)
    """
).set(observed=data)
density = model.logdensity("posterior")

a0, b0, k0, m0, n = 3.0, 2.0, 2.0, 0.0, data.size
mean = data.mean()
kn, an = k0 + n, a0 + n / 2
mn = (k0 * m0 + n * mean) / kn
bn = b0 + 0.5 * ((data - mean) ** 2).sum() + k0 * n * (mean - m0) ** 2 / (2 * kn)
exact = {"mu": (mn, bn / (kn * (an - 1))), "tau": (an / bn, an / bn**2)}
log_evidence = (
    lgamma(an) - lgamma(a0) + a0 * log(b0) - an * log(bn)
    + 0.5 * log(k0 / kn) - n / 2 * log(2 * pi)
)  # fmt: skip

Check the normalization

The free coordinates’ density integrates to the closed-form evidence.

mu_axis = jnp.linspace(-3.0, 5.0, 401)
log_tau_axis = jnp.linspace(-4.0, 3.0, 401)
grid = jax.vmap(
    lambda m: jax.vmap(lambda t: density({"mu": m, "tau": t}))(log_tau_axis)
)(mu_axis)
total = np.trapezoid(
    np.trapezoid(np.exp(np.asarray(grid)), np.asarray(log_tau_axis), axis=1),
    np.asarray(mu_axis),
)
assert abs(np.log(total) - log_evidence) < 1e-6

Draw from four chains

BlackJAX NUTS with window adaptation, as in Sample and approximate a posterior with BlackJAX. The sampler statistics ride along for ArviZ.

chains, warmup, draws = 4, 400, 1000
starts = jax.vmap(
    lambda key: jax.tree.map(
        lambda x: x + 0.5 * jax.random.normal(key, jnp.shape(x)), density.init
    )
)(jax.random.split(jax.random.key(1), chains))


def run_chain(key: jax.Array, start: Any) -> Any:
    """Adapt one chain, then keep its draws and sampler statistics."""
    warmup_key, sample_key = jax.random.split(key)
    adaptation = blackjax.window_adaptation(blackjax.nuts, density.fn)
    (state, parameters), _ = adaptation.run(warmup_key, start, num_steps=warmup)
    kernel = blackjax.nuts(density.fn, **parameters)

    def step(state: Any, key: jax.Array) -> tuple[Any, Any]:
        state, info = kernel.step(key, state)
        stats = {"diverging": info.is_divergent, "energy": info.energy}
        return state, (state.position, stats)

    return jax.lax.scan(step, state, jax.random.split(sample_key, draws))[1]


free, stats = jax.jit(jax.vmap(run_chain))(
    jax.random.split(jax.random.key(2), chains), starts
)
samples = jax.vmap(jax.vmap(density.constrain))(free)

Build the ArviZ data

One variable per name in density.names. A parameter of shape density.shapes[name] keeps those trailing dimensions after (chain, draw); both parameters here are scalars.

posterior = az.from_dict(
    {
        "posterior": {name: np.asarray(samples[name]) for name in density.names},
        "sample_stats": {name: np.asarray(value) for name, value in stats.items()},
    }
)
summary = az.summary(posterior, ci_prob=0.9, round_to="none")
print(summary.to_string())

         mean        sd  eti90_lb  eti90_ub     ess_bulk     ess_tail     r_hat  mcse_mean   mcse_sd
mu   0.919051  0.268846  0.475752  1.359296  2460.121136  2235.982077  0.999856   0.005433  0.004612
tau  1.600658  0.612511  0.745355  2.748949  3012.758843  2659.248743  1.000553   0.011030  0.010007

Check the summary against the closed form

Each posterior mean lies within four of ArviZ’s Monte Carlo standard errors of the exact value, and each standard deviation within five percent. Rank normalized R-hat and bulk ESS confirm the chains mixed.

for name, (exact_mean, exact_variance) in exact.items():
    row = {key: float(value) for key, value in summary.loc[name].items()}
    assert abs(row["mean"] - exact_mean) < 4 * row["mcse_mean"], name
    assert abs(row["sd"] / np.sqrt(exact_variance) - 1) < 0.05, name
    assert row["r_hat"] < 1.01 and row["ess_bulk"] > 400, name
assert not np.asarray(stats["diverging"]).any()