Sample a posterior with NumPyro’s NUTS

NumPyro’s MCMC runs NUTS on any potential, the negated log density. flatppl.Module.logdensity() supplies it in unconstrained coordinates together with a starting point and the map back to the model’s parameters. The model is a Normal likelihood with unknown mean and precision under its conjugate Normal-Gamma prior, so the posterior and its evidence are known exactly.

NumPyro’s variational inference fits guides to a NumPyro model, so a foreign log density cannot reach it; Sample and approximate a posterior with BlackJAX covers variational inference on the same density.

Runtime budget: 120 s.

from math import lgamma, log, pi

import jax
import jax.numpy as jnp
import numpy as np
from numpyro.diagnostics import effective_sample_size, split_gelman_rubin
from numpyro.infer import MCMC, NUTS

from flatppl import flatppl

The model and its closed-form posterior

tau | y ~ Gamma(a_n, b_n), and 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 density of the free coordinates (mu, log tau) integrates to the evidence, so the potential carries every constant.

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),
)
print(f"log evidence: grid {np.log(total):.6f}, closed form {log_evidence:.6f}")
assert abs(np.log(total) - log_evidence) < 1e-6

log evidence: grid -10.558731, closed form -10.558731

Run NUTS

potential_fn takes the position pytree. With several chains, init_params carries a leading chain axis. "vectorized" runs the chains with jax.vmap on one device.

chains = 4
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))
mcmc = MCMC(
    NUTS(potential_fn=density.potential),
    num_warmup=400,
    num_samples=1000,
    num_chains=chains,
    chain_method="vectorized",
    progress_bar=False,
)
mcmc.run(jax.random.key(2), init_params=starts)
free = mcmc.get_samples(group_by_chain=True)
samples = jax.vmap(jax.vmap(density.constrain))(free)

Compare with the closed form

Each mean must lie within four Monte Carlo standard errors of the exact value.

for name, (exact_mean, exact_variance) in exact.items():
    chain_draws = np.asarray(samples[name])
    ess = float(effective_sample_size(chain_draws))
    rhat = float(split_gelman_rubin(chain_draws))
    estimate = float(chain_draws.mean())
    error = np.sqrt(exact_variance / ess)
    print(
        f"{name}: NUTS {estimate:.4f}, exact {exact_mean:.4f}, "
        f"MCSE {error:.4f}, ESS {ess:.0f}, R-hat {rhat:.3f}"
    )
    assert abs(estimate - exact_mean) < 4 * error
    assert rhat < 1.01

mu: NUTS 0.9246, exact 0.9200, MCSE 0.0060, ESS 2042, R-hat 1.001
tau: NUTS 1.5962, exact 1.5916, MCSE 0.0126, ESS 2281, R-hat 1.001