MCMC diagnostics + ensemble uncertainty: plots, EnsembleEnergy, disagreement OOD

[davidkhjo]

Two threads that make EBMs easier to trust: is the sampler mixing? and is this input in-distribution?

Diagnostics

• eval.autocorrelation(samples, max_lag) — mean per-chain ACF ρ̂_t (FFT, the same autocovariance as effective_sample_size).
• viz.autocorrelation_plot (with a ±1.96/√N band), viz.rank_plot (Vehtari rank histograms — uniform under good mixing), viz.trace_plot.

Ensemble uncertainty

• EnsembleEnergy — a deep-ensemble mean energy E = (1/N)ΣE_i (the geometric-mean density; equals SumEnergy with weights 1/N), deliberately distinct from MixtureEnergy's logsumexp. Adds member_energies and disagreement (per-sample variance across members).
• eval.ensemble_disagreement — that variance, batched, as an epistemic OOD score.

Validation (closed-form / distributional)

• ACF estimator matches an AR(1)'s ρ^t (to N-error).
• Ensemble of Gaussian-energy members reproduces the closed-form combined Gaussian variance 1/mean(1/σ_i²) exactly.
• Member disagreement separates in-distribution from OOD at AUROC ≈ 1; rank plots are uniform for matched chains and skewed for an offset chain.

Notes

• Example ensemble_ood.py: a 3-member ensemble on two-moons; disagreement flags a Gaussian blob at AUROC ≈ 1.
• torch-only; public API preserved; accumulates under the unreleased 0.15.0. This completes the E1–E6 roadmap; NUTS and the latent/hierarchical EBM (E7) are left as follow-ups (NUTS deferred because a correct batched implementation is easy to bias silently — against this library's correctness-first stance).

🤖 Generated with Claude Code