MCMC fit check

What it does

• Reads an InferenceData and checks divergences, R-hat, bulk and tail effective sample size, BFMI per chain and tree-depth saturation, with the session's thresholds.
• Prints a verdict: the computation is sound and the predictions are next, or do not interpret this posterior yet, naming the checks that failed.
• A second function draws one posterior predictive check, sampling the predictive group first if the model has not been asked for it yet.
• Run as a script it fits a well-specified model and a centered funnel, so you can see what a pass and a fail look like before pointing it at your own model.

How to use it

Drop the file beside your notebook, import it, and call check_fit on the result of pm.sample before you read a single parameter estimate.

Needs: arviz, matplotlib; PyMC for the self-test. License: MIT. Written from the method the session teaches; the instructor's own materials come with the recording.

The session it came from

MCMC in Practice. The skill is the method from that session, packaged so you can use it before you have watched it; the session is where the reasoning lives. Every skill in the library names its session the same way.

The file

mcmc_checks.py · 119 lines

"""mcmc_checks: the fit-checking pass from the AthlyticZ session "MCMC in Practice" (Chris Fonnesbeck),
as one reusable function you can call on anything pm.sample() returns.

    import mcmc_checks as mc
    report = mc.check_fit(idata)          # prints the verdict, returns a dict
    mc.ppc_overlay(idata, "y")            # one posterior predictive check, as a plot

What it checks, in the order the session teaches it:
  1. divergences            any at all means the sampler hit geometry it could not follow (think non-centered)
  2. R-hat                  > 1.01 on any parameter: the chains disagree, do not trust the posterior yet
  3. effective sample size  bulk or tail ESS under 400: the estimate is noisier than the draw count suggests
  4. BFMI                   under 0.3 on any chain: the energy transitions are poor, often a funnel
  5. tree depth             hits at the cap mean NUTS was cut short; raise the cap or reparameterize

Requires: arviz >= 0.17, matplotlib. Works on InferenceData from PyMC 5 and 6 (nutpie included).
License: MIT. Written for athlyticz.com/skills from the method the session teaches; the session's own
notebook and datasets come with the recording.
"""
from __future__ import annotations

import numpy as np

try:
    import arviz as az
except ImportError as e:  # pragma: no cover
    raise SystemExit("mcmc_checks needs arviz: pip install arviz") from e


def check_fit(idata, var_names=None, rhat_max=1.01, ess_min=400, bfmi_min=0.3, verbose=True):
    """Run the five sampler checks on an InferenceData and return {check: (ok, detail)}.

    var_names narrows the R-hat and ESS tables to the parameters you care about (defaults to all).
    The thresholds are the session's: R-hat 1.01, ESS 400, BFMI 0.3.
    """
    out = {}
    ss = idata.get("sample_stats")

    # 1. divergences
    if ss is not None and "diverging" in ss:
        n_div = int(ss["diverging"].sum())
        out["divergences"] = (n_div == 0, f"{n_div} divergent transition(s)")
    else:
        out["divergences"] = (True, "no sample_stats.diverging recorded")

    # 2 and 3. R-hat and ESS, over the posterior
    summ = az.summary(idata, var_names=var_names, kind="diagnostics")
    worst_rhat = float(summ["r_hat"].max())
    out["r_hat"] = (worst_rhat <= rhat_max, f"max R-hat {worst_rhat:.3f} ({summ['r_hat'].idxmax()})")
    ess_bulk, ess_tail = float(summ["ess_bulk"].min()), float(summ["ess_tail"].min())
    out["ess"] = (min(ess_bulk, ess_tail) >= ess_min,
                  f"min bulk ESS {ess_bulk:.0f} ({summ['ess_bulk'].idxmin()}), min tail ESS {ess_tail:.0f} ({summ['ess_tail'].idxmin()})")

    # 4. BFMI, per chain
    try:
        bfmi = az.bfmi(idata)
        out["bfmi"] = (bool(np.all(bfmi >= bfmi_min)), "per-chain BFMI " + ", ".join(f"{b:.2f}" for b in bfmi))
    except Exception as e:  # no energy recorded (e.g. a non-HMC sampler)
        out["bfmi"] = (True, f"BFMI not available ({type(e).__name__})")

    # 5. tree depth saturation
    if ss is not None and "tree_depth" in ss:
        depth = ss["tree_depth"].values
        cap = int(depth.max())
        frac = float((depth == cap).mean())
        out["tree_depth"] = (frac < 0.05 or cap < 10, f"max depth {cap}, {frac:.1%} of transitions at it")
    else:
        out["tree_depth"] = (True, "no tree_depth recorded")

    if verbose:
        width = max(len(k) for k in out)
        for k, (ok, detail) in out.items():
            print(f"{'ok  ' if ok else 'FAIL'} {k:<{width}}  {detail}")
        bad = [k for k, (ok, _) in out.items() if not ok]
        print("verdict:", "the computation is sound; now check the predictions" if not bad
              else "do not interpret this posterior yet: " + ", ".join(bad))
    return out


def ppc_overlay(idata, var, model=None, draws=100, ax=None):
    """One posterior predictive check: simulated datasets over the observed one.

    If idata has no posterior_predictive group, pass the PyMC model and it is sampled here.
    Returns the matplotlib Axes. The question to ask of the picture is the session's: do the
    simulated datasets look like the real one, or is the model missing something the
    sampler diagnostics could never show?
    """
    import matplotlib.pyplot as plt

    if "posterior_predictive" not in idata:
        if model is None:
            raise ValueError("no posterior_predictive group: pass model= so it can be sampled")
        import pymc as pm
        with model:
            pm.sample_posterior_predictive(idata, extend_inferencedata=True, random_seed=1)
    ax = az.plot_ppc(idata, var_names=[var], num_pp_samples=draws, ax=ax)
    ax.set_title(f"posterior predictive check: {var}")
    return ax


if __name__ == "__main__":
    # a two-minute self-test: a well-specified model, then a funnel that should fail
    import pymc as pm

    rng = np.random.default_rng(1)
    y = rng.normal(2.0, 1.0, size=200)
    with pm.Model() as good:
        mu = pm.Normal("mu", 0, 5)
        sd = pm.HalfNormal("sd", 2)
        pm.Normal("y", mu, sd, observed=y)
        idata = pm.sample(1000, tune=1000, chains=4, progressbar=False, random_seed=1)
    print("well-specified model:")
    check_fit(idata)

    with pm.Model() as funnel:
        tau = pm.Normal("tau", 0, 3)
        x = pm.Normal("x", 0, pm.math.exp(tau), shape=9)
        idata_f = pm.sample(1000, tune=1000, chains=4, progressbar=False, random_seed=1)
    print("\nNeal's funnel (centered):")
    check_fit(idata_f)