diff --git a/CHANGELOG.md b/CHANGELOG.md index e64ae77..5910252 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,14 @@ Internal cleanup plus a diffusion track, exact-likelihood flows, and conditional-generation / calibration tools. +- **Diagnostics & ensemble uncertainty.** `viz.autocorrelation_plot` / + `rank_plot` / `trace_plot` and a numeric `eval.autocorrelation` for judging MCMC + mixing; `EnsembleEnergy` (a deep-ensemble *mean* energy — the geometric-mean + density, distinct from `MixtureEnergy`'s logsumexp) with `member_energies` / + `disagreement`, plus `eval.ensemble_disagreement` as an epistemic OOD score. + Validated: the ACF estimator matches an AR(1)'s `ρ^t`; the ensemble reproduces + the closed-form combined Gaussian; disagreement separates in-distribution from + OOD at AUROC ≈ 1. Example `ensemble_ood.py`. - **`nets.NeuralSplineCouplingFlow`** — a rational-quadratic neural spline flow (Durkan et al. 2019): monotonic-spline coupling layers dropped into the affine-flow scaffold, strictly more expressive per layer, so it fits sharp diff --git a/README.md b/README.md index b355e3c..5e925e9 100644 --- a/README.md +++ b/README.md @@ -47,10 +47,10 @@ returns `LossOutput(loss, metrics, x_neg)`; call `out.loss.backward()`). | **Energies** | any callable `(B, *shape) -> (B,)`; `nets.MLPEnergy` / `ConvEnergy` / `ResNetEnergy` / `ConvClassifier` (SiLU, optional spectral norm, no batch norm), `nets.RBM` (Bernoulli RBM with exact `log_z`), `GuidedEnergy` (classifier-free guidance), `nets.IsingEnergy` / `PottsEnergy` (discrete lattices), `nets.FunnelEnergy` / `GaussianMixtureEnergy` / `BananaEnergy` (closed-form targets), `nets.AffineCouplingFlow` (RealNVP) / `nets.NeuralSplineCouplingFlow` (rational-quadratic spline) / `nets.ContinuousNormalizingFlow` (FFJORD — trainable exact-likelihood flows / self-normalized energies), noise-conditional variants for NCSN; `EnergyModel`, `ebm.score` | | **Samplers** | `LangevinDynamics` (ULA/SGLD), `MALA`, `AdaptiveMALA` (dual-averaging step-size warmup + diagonal metric), `HMC`, `UnderdampedLangevin` (SGHMC), `PreconditionedLangevin`, `ParallelTempering` (replica exchange), `TemperedTransitions`, `SVGD` (Stein variational), `GibbsSampler` (block Gibbs), `GibbsWithGradients` + `CategoricalGibbsWithGradients`, `AnnealedLangevinDynamics`, `ProbabilityFlowODE` / `PredictorCorrector` (score-SDE), `DDPMAncestralSampler` (VP diffusion) | | **Losses** | `ContrastiveDivergence` (CD-k / persistent CD), `DiffusionRecoveryLikelihood` + `drl_sample`, `DenoisingScoreMatching` / `MultiSigmaDenoisingScoreMatching` (NCSN), `VPDenoisingScoreMatching` (DDPM), `SlicedScoreMatching`, `ExactScoreMatching`, `EnergyDiscrepancy` (MCMC-free), `PseudoLikelihood` / `RatioMatching` / `ConcreteScoreMatching` (MCMC-free, discrete), `NoiseContrastiveEstimation`, `JEMLoss` | -| **Composition** | `SumEnergy` (product of experts), `MixtureEnergy`, `TemperedEnergy` — energies compose like densities and nest | +| **Composition** | `SumEnergy` (product of experts), `MixtureEnergy`, `EnsembleEnergy` (deep-ensemble mean energy + member disagreement), `TemperedEnergy` — energies compose like densities and nest | | **Training** | thin `Trainer` (device, EMA, supervised batches, `save`/`load` checkpointing), `ReplayBuffer`, `EMA` | -| **Eval** | `ais_log_z` / `reverse_ais_log_z` (bracket `log Z`), `pf_ode_log_likelihood` (exact likelihood via the probability-flow ODE), `bits_per_dim`, `frechet_distance` (FID), `mmd`, `precision_recall`, `inception_score`, `kernel_stein_discrepancy` / `classifier_two_sample_test` / `fisher_divergence` (goodness-of-fit), `mutual_information` (MINE), `expected_calibration_error` / `reliability_curve` / `temperature_scale` (calibration), `ood_auroc`, `effective_sample_size` / `split_rhat` (MCMC diagnostics) | -| **Data & viz** | 2D toys (`two_moons`, `eight_gaussians`, `checkerboard`, `rings`, `spirals`) and torchvision-free image loaders (`mnist`, `fashion_mnist`, `cifar10`, `cifar100`); `viz.energy_contour` / `plot_samples` / `energy_histogram` / `show_images` | +| **Eval** | `ais_log_z` / `reverse_ais_log_z` (bracket `log Z`), `pf_ode_log_likelihood` (exact likelihood via the probability-flow ODE), `bits_per_dim`, `frechet_distance` (FID), `mmd`, `precision_recall`, `inception_score`, `kernel_stein_discrepancy` / `classifier_two_sample_test` / `fisher_divergence` (goodness-of-fit), `mutual_information` (MINE), `expected_calibration_error` / `reliability_curve` / `temperature_scale` (calibration), `ood_auroc`, `ensemble_disagreement` (epistemic OOD), `effective_sample_size` / `split_rhat` / `autocorrelation` (MCMC diagnostics) | +| **Data & viz** | 2D toys (`two_moons`, `eight_gaussians`, `checkerboard`, `rings`, `spirals`) and torchvision-free image loaders (`mnist`, `fashion_mnist`, `cifar10`, `cifar100`); `viz.energy_contour` / `plot_samples` / `energy_histogram` / `show_images` / `autocorrelation_plot` / `rank_plot` / `trace_plot` | ## Examples @@ -77,6 +77,7 @@ Runnable scripts in [`examples/`](https://github.com/davidkhjo/ebmkit/tree/main/ - `sampling_hard_targets.py` — parallel tempering escapes a trapped mode; ESS / R̂ diagnostics - `adaptive_mala.py` — a self-tuning MALA: dual-averaging step size + a learned diagonal metric - `goodness_of_fit.py` — KSD for model selection; classifier two-sample test +- `ensemble_ood.py` — a deep-ensemble EBM whose member disagreement flags OOD - `mine_mutual_information.py` — estimate mutual information with MINE vs the Gaussian closed form - `benchmark_samplers.py` — rank samplers on the banana against exact draws (ESS, R̂, MMD) - `checkpoint_resume.py` — save a run and resume it in a fresh process diff --git a/examples/ensemble_ood.py b/examples/ensemble_ood.py new file mode 100644 index 0000000..a38dee5 --- /dev/null +++ b/examples/ensemble_ood.py @@ -0,0 +1,76 @@ +"""Deep-ensemble EBM: member disagreement as an epistemic OOD signal. + +A single energy can be confidently wrong off-distribution. An ensemble of energy +networks trained on the same data agrees where it saw data and *disagrees* where +it didn't — so the variance of the member energies is an epistemic-uncertainty +score that flags OOD inputs. `EnsembleEnergy` pools the members into a mean energy +(a geometric-mean density) for sampling/scoring, and `ensemble_disagreement` +returns the per-sample variance; here it separates two-moons (in-distribution) +from a Gaussian blob (OOD) at AUROC ≈ 1. + +Run: python examples/ensemble_ood.py +Outputs ensemble_ood_result.png next to this script (needs the [viz] extra). +""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +import ebm + + +def _train_member(data: torch.Tensor, seed: int) -> torch.nn.Module: + torch.manual_seed(seed) + net = ebm.nets.MLPEnergy(dim=2, hidden=(128, 128)) + sampler = ebm.LangevinDynamics(step_size=0.01, steps=60) + loss_fn = ebm.ContrastiveDivergence( + sampler, buffer=ebm.ReplayBuffer(8192, (2,)), energy_reg=0.1 + ) + trainer = ebm.Trainer(net, loss_fn, lr=1e-3, ema_decay=0.999, device="cpu") + trainer.fit(data, steps=3000, batch_size=256) + return net + + +def main() -> None: + torch.manual_seed(0) + data = ebm.datasets.two_moons(8192) + ensemble = ebm.EnsembleEnergy(*[_train_member(data, seed) for seed in range(3)]) + + x_in = ebm.datasets.two_moons(2000) + x_out = torch.randn(2000, 2) * 0.6 + torch.tensor([0.0, 3.0]) # off-manifold blob + d_in = ebm.eval.ensemble_disagreement(ensemble, x_in) + d_out = ebm.eval.ensemble_disagreement(ensemble, x_out) + auroc = ebm.eval.ood_auroc(lambda z: ebm.eval.ensemble_disagreement(ensemble, z), x_in, x_out) + print(f"mean disagreement in-dist {d_in.mean():.3f} OOD {d_out.mean():.3f}") + print(f"OOD AUROC (disagreement) = {auroc:.3f}") + + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, axes = plt.subplots(1, 2, figsize=(12, 5)) + grid = torch.linspace(-3, 4, 200) + gy, gx = torch.meshgrid(grid, grid, indexing="ij") + pts = torch.stack([gx.reshape(-1), gy.reshape(-1)], dim=1) + dis = ebm.eval.ensemble_disagreement(ensemble, pts).reshape(200, 200) + im = axes[0].imshow(dis, extent=(-3, 4, -3, 4), origin="lower", vmax=float(dis.quantile(0.98))) + axes[0].scatter(x_in[:, 0], x_in[:, 1], s=3, color="white", alpha=0.3) + axes[0].set_title("member disagreement (bright = uncertain / OOD)") + fig.colorbar(im, ax=axes[0], fraction=0.046) + + axes[1].hist(d_in.numpy(), bins=50, alpha=0.6, label="two-moons (in)", density=True) + axes[1].hist(d_out.numpy(), bins=50, alpha=0.6, label="Gaussian blob (OOD)", density=True) + axes[1].set_xlabel("ensemble disagreement") + axes[1].set_title(f"OOD AUROC = {auroc:.3f}") + axes[1].legend() + + out = Path(__file__).parent / "ensemble_ood_result.png" + fig.savefig(out, dpi=120, bbox_inches="tight") + print(f"saved {out}") + + +if __name__ == "__main__": + main() diff --git a/examples/ensemble_ood_result.png b/examples/ensemble_ood_result.png new file mode 100644 index 0000000..5459e7b Binary files /dev/null and b/examples/ensemble_ood_result.png differ diff --git a/src/ebm/__init__.py b/src/ebm/__init__.py index c910478..af45d6f 100644 --- a/src/ebm/__init__.py +++ b/src/ebm/__init__.py @@ -7,7 +7,7 @@ from ebm import datasets, eval, nets, viz from ebm.ais import AISResult, ais_log_z, log_likelihood, reverse_ais_log_z from ebm.buffer import ReplayBuffer -from ebm.compose import MixtureEnergy, SumEnergy, TemperedEnergy +from ebm.compose import EnsembleEnergy, MixtureEnergy, SumEnergy, TemperedEnergy from ebm.diffusion import DDPMAncestralSampler, VPDenoisingScoreMatching, VPSchedule from ebm.energy import ConditionalEnergyFn, EnergyFn, EnergyModel, score from ebm.jem import ClassifierEnergy, ConditionalEnergy, GuidedEnergy, JEMLoss @@ -70,6 +70,7 @@ "EnergyDiscrepancy", "EnergyFn", "EnergyModel", + "EnsembleEnergy", "ExactScoreMatching", "GibbsSampler", "GibbsWithGradients", diff --git a/src/ebm/compose.py b/src/ebm/compose.py index ca2ab47..6941428 100644 --- a/src/ebm/compose.py +++ b/src/ebm/compose.py @@ -82,6 +82,36 @@ def forward(self, x: Tensor) -> Tensor: return -torch.logsumexp(stacked, dim=0) +class EnsembleEnergy(_Composite): + """Deep ensemble as a **mean energy**: ``E(x) = (1/N) Σ_i E_i(x)``. + + Averaging energies is the *geometric* mean of the members' densities + (``p ∝ Π p_i^{1/N}``) — the variance-reduced log-density estimate you want + from a deep ensemble, and equivalent to ``SumEnergy`` with weights ``1/N``. + This is deliberately distinct from `MixtureEnergy`, which mixes densities + with a ``logsumexp`` (an *arithmetic* mean). + + Beyond the pooled energy, the spread of the members carries epistemic + uncertainty: `member_energies(x)` returns the ``(B, N)`` per-member energies + and `disagreement(x)` their per-sample variance — low where the members agree + (seen data), high off-distribution. See `eval.ensemble_disagreement`. + """ + + def __init__(self, *energies: EnergyFn): + super().__init__(energies) + + def member_energies(self, x: Tensor) -> Tensor: + """Per-member energies, stacked as ``(B, N)``.""" + return torch.stack([e(x) for e in self.energies], dim=1) + + def forward(self, x: Tensor) -> Tensor: + return self.member_energies(x).mean(dim=1) + + def disagreement(self, x: Tensor) -> Tensor: + """Per-sample variance across members ``(B,)`` — an epistemic OOD signal.""" + return self.member_energies(x).var(dim=1, unbiased=False) + + class TemperedEnergy(nn.Module): """``E(x) / T``: temperature ``T > 1`` flattens ``p``, ``T < 1`` sharpens it. diff --git a/src/ebm/eval/__init__.py b/src/ebm/eval/__init__.py index 5ae16d0..e5d365e 100644 --- a/src/ebm/eval/__init__.py +++ b/src/ebm/eval/__init__.py @@ -12,8 +12,10 @@ temperature_scale, ) from ebm.eval.diagnostics import ( + autocorrelation, effective_sample_size, energies, + ensemble_disagreement, ood_auroc, split_rhat, ) @@ -34,10 +36,12 @@ __all__ = [ "AISResult", "ais_log_z", + "autocorrelation", "bits_per_dim", "classifier_two_sample_test", "effective_sample_size", "energies", + "ensemble_disagreement", "expected_calibration_error", "fisher_divergence", "frechet_distance", diff --git a/src/ebm/eval/diagnostics.py b/src/ebm/eval/diagnostics.py index 46301c9..5edc095 100644 --- a/src/ebm/eval/diagnostics.py +++ b/src/ebm/eval/diagnostics.py @@ -2,11 +2,16 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import torch from torch import Tensor from ebm.energy import EnergyFn +if TYPE_CHECKING: + from ebm.compose import EnsembleEnergy + def _as_chains(samples: Tensor) -> Tensor: """Coerce MCMC output to ``(n_chains, n_samples, dim)`` in float64.""" @@ -104,6 +109,28 @@ def effective_sample_size(samples: Tensor) -> Tensor: return ess +@torch.no_grad() +def autocorrelation(samples: Tensor, max_lag: int = 40) -> Tensor: + """Mean per-chain autocorrelation ``ρ̂_t`` for lags ``0..max_lag``, per dimension. + + Input is ``(n_chains, n_samples[, dim])``. Returns ``(max_lag+1, dim)`` (lag 0 + is ``1``). Computed by FFT (zero-padded past ``2N``) and averaged over chains — + the same autocovariance used by `effective_sample_size`. A chain that mixes + well decays toward zero within a few lags; slow decay flags autocorrelation. + """ + x = _as_chains(samples) + _, n, _ = x.shape + max_lag = min(max_lag, n - 1) + centered = x - x.mean(dim=1, keepdim=True) + n_fft = 1 + while n_fft < 2 * n: + n_fft <<= 1 + spec = torch.fft.rfft(centered, n=n_fft, dim=1) + acov = torch.fft.irfft(spec.abs().pow(2), n=n_fft, dim=1)[:, :n, :] # (M, N, d) + acf = acov / acov[:, :1, :] # normalize by lag-0 per chain/dim + return acf.mean(dim=0)[: max_lag + 1] + + @torch.no_grad() def energies(energy: EnergyFn, x: Tensor, batch_size: int = 1024) -> Tensor: """Energies of ``x`` computed in batches, returned on CPU.""" @@ -111,6 +138,20 @@ def energies(energy: EnergyFn, x: Tensor, batch_size: int = 1024) -> Tensor: return torch.cat(out) +@torch.no_grad() +def ensemble_disagreement(ensemble: EnsembleEnergy, x: Tensor, batch_size: int = 1024) -> Tensor: + """Per-sample ensemble disagreement (variance of member energies), batched. + + The spread of an ensemble's member energies is an *epistemic* uncertainty + signal: members trained on the same data agree where they saw data and + diverge off-distribution, so `ensemble_disagreement` tends to be much larger + on OOD inputs — feed it to `ood_auroc` as a detector. Returns a ``(B,)`` tensor + on CPU. + """ + out = [ensemble.disagreement(chunk).detach().cpu() for chunk in x.split(batch_size)] + return torch.cat(out) + + @torch.no_grad() def ood_auroc(energy: EnergyFn, x_in: Tensor, x_out: Tensor) -> float: """AUROC for separating in-distribution from OOD data by energy. diff --git a/src/ebm/viz.py b/src/ebm/viz.py index b2cc9b2..569d441 100644 --- a/src/ebm/viz.py +++ b/src/ebm/viz.py @@ -9,7 +9,15 @@ from ebm.energy import EnergyFn -__all__ = ["energy_contour", "energy_histogram", "plot_samples", "show_images"] +__all__ = [ + "autocorrelation_plot", + "energy_contour", + "energy_histogram", + "plot_samples", + "rank_plot", + "show_images", + "trace_plot", +] def _require_matplotlib(): @@ -114,6 +122,86 @@ def show_images( return ax +@torch.no_grad() +def _chain_ranks(samples: Tensor, dim: int = 0) -> Tensor: + """Rank each draw among the pooled draws of all chains; returns ``(n_chains, n_samples)``.""" + from ebm.eval.diagnostics import _as_chains + + x = _as_chains(samples)[..., dim] # (M, N) + m, n = x.shape + order = x.reshape(-1).argsort() + ranks = torch.empty(m * n, dtype=torch.float64) + ranks[order] = torch.arange(1, m * n + 1, dtype=torch.float64) + return ranks.reshape(m, n) + + +@torch.no_grad() +def autocorrelation_plot(samples: Tensor, max_lag: int = 40, dim: int = 0, ax=None): + """Autocorrelation ``ρ̂_t`` vs lag for MCMC output, with a ``±1.96/√N`` band. Returns the axes. + + Input is ``(n_chains, n_samples[, dim])``. Bars decaying to inside the band + within a few lags indicate good mixing; a slow decay flags autocorrelation. + """ + from ebm.eval.diagnostics import autocorrelation + + plt = _require_matplotlib() + acf = autocorrelation(samples, max_lag)[:, dim] + n = samples.shape[1] + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + ax.bar(range(len(acf)), acf.numpy(), width=0.8, color="#5c50c9") + band = 1.96 / (n**0.5) + for y in (band, -band): + ax.axhline(y, ls="--", color="gray", lw=1) + ax.axhline(0.0, color="black", lw=0.8) + ax.set_xlabel("lag") + ax.set_ylabel("autocorrelation") + return ax + + +@torch.no_grad() +def rank_plot(samples: Tensor, dim: int = 0, bins: int = 20, ax=None): + """Vehtari rank histogram per chain (uniform under good mixing). Returns the axes. + + Ranks every draw among the pooled draws of all chains, then histograms each + chain's ranks. Well-mixed chains give flat, overlapping histograms near the + dashed uniform line; a chain offset from the others shows a sloped or skewed + histogram. Reference: Vehtari et al. (2021). + """ + plt = _require_matplotlib() + ranks = _chain_ranks(samples, dim) + m, n = ranks.shape + edges = torch.linspace(1, m * n, bins + 1).numpy() + if ax is None: + _, ax = plt.subplots(figsize=(6, 4)) + for i in range(m): + ax.hist(ranks[i].numpy(), bins=edges, histtype="step", alpha=0.7) + ax.axhline(n / bins, ls="--", color="gray", lw=1) # uniform expectation per bin + ax.set_xlabel("rank") + ax.set_ylabel("count") + return ax + + +@torch.no_grad() +def trace_plot(samples: Tensor, dim: int = 0, max_chains: int = 8, ax=None): + """Trace of each chain's value over iterations (mixing at a glance). Returns the axes. + + Input is ``(n_chains, n_samples[, dim])``. Chains that overlap and wander + across the same range have mixed; a chain stuck at a different level has not. + """ + from ebm.eval.diagnostics import _as_chains + + plt = _require_matplotlib() + x = _as_chains(samples)[..., dim] + if ax is None: + _, ax = plt.subplots(figsize=(8, 4)) + for i in range(min(x.shape[0], max_chains)): + ax.plot(x[i].numpy(), alpha=0.7, lw=0.8) + ax.set_xlabel("iteration") + ax.set_ylabel(f"dim {dim}") + return ax + + def energy_histogram(energy: EnergyFn, batches: dict[str, Tensor], ax=None, bins: int = 60): """Overlaid energy histograms, e.g. ``{"data": x, "samples": x_neg}``. diff --git a/tests/test_compose.py b/tests/test_compose.py index 091a3bf..10b9255 100644 --- a/tests/test_compose.py +++ b/tests/test_compose.py @@ -77,3 +77,32 @@ def test_compositions_nest(): ebm.SumEnergy(quadratic_energy, ebm.MixtureEnergy(quadratic_energy)), 2.0 ) assert nested(torch.randn(8, 2)).shape == (8,) + + +def test_ensemble_energy_is_geometric_mean_gaussian(): + # Members N(0, σ_i²) → mean energy is N(0, 1/precision), precision = mean(1/σ_i²). + sigmas = [1.0, 0.5, 2.0] + members = [gaussian_energy(0.0, s**2) for s in sigmas] + ens = ebm.EnsembleEnergy(*members) + precision = sum(1 / s**2 for s in sigmas) / len(sigmas) + samples = ebm.MALA(step_size=0.1, steps=800).sample(ens, torch.randn(6000, 2)) + assert (samples.var(0) - 1 / precision).abs().max().item() < 0.05 + # member_energies stacks per-member; forward is their mean. + x = torch.randn(16, 2) + me = ens.member_energies(x) + assert me.shape == (16, 3) + assert torch.allclose(ens(x), me.mean(dim=1)) + + +def test_ensemble_disagreement_flags_ood(): + sigmas = [1.0, 0.5, 2.0] + ens = ebm.EnsembleEnergy(*[gaussian_energy(0.0, s**2) for s in sigmas]) + x_in = torch.randn(2000, 2) + x_out = 6.0 + torch.randn(2000, 2) + d_in = ebm.eval.ensemble_disagreement(ens, x_in) + d_out = ebm.eval.ensemble_disagreement(ens, x_out) + assert d_in.shape == (2000,) + assert d_out.mean() > 50 * d_in.mean() # members diverge off-distribution + # As an OOD score it separates the two sets almost perfectly. + auroc = ebm.eval.ood_auroc(lambda z: ebm.eval.ensemble_disagreement(ens, z), x_in, x_out) + assert auroc > 0.99 diff --git a/tests/test_datasets_eval.py b/tests/test_datasets_eval.py index e3625e6..037cacd 100644 --- a/tests/test_datasets_eval.py +++ b/tests/test_datasets_eval.py @@ -47,6 +47,16 @@ def test_ess_matches_ar1_closed_form(rho): assert abs(ess - predicted) / predicted < 0.2 +def test_autocorrelation_matches_ar1(): + # A stationary AR(1) has ρ̂_t = ρ^t; check the estimator recovers it. + m, n, rho = 16, 4000, 0.7 + acf = ebm.eval.autocorrelation(_ar1(m, n, rho), max_lag=8) + assert acf.shape == (9, 1) + theory = torch.tensor([rho**k for k in range(9)]) + assert (acf[:, 0] - theory).abs().max().item() < 0.03 + assert acf[0, 0].item() == pytest.approx(1.0) + + def test_diagnostics_reject_too_short(): with pytest.raises(ValueError): ebm.eval.effective_sample_size(torch.randn(4, 3, 2)) diff --git a/tests/test_viz.py b/tests/test_viz.py index ca48f49..aa8d6c6 100644 --- a/tests/test_viz.py +++ b/tests/test_viz.py @@ -54,3 +54,31 @@ def test_energy_contour_plot_samples_and_histogram(): ) assert ax3.get_xlabel() == "energy" assert len(ax3.get_legend().get_texts()) == 2 + + +def test_mcmc_diagnostic_plots_smoke_and_rank_uniformity(): + import matplotlib + + matplotlib.use("Agg") + + # Well-mixed chains: N(0,1) i.i.d.; one offset chain to skew the rank plot. + good = torch.randn(6, 500, 2) + ax = ebm.viz.autocorrelation_plot(good, max_lag=20, dim=0) + assert len(ax.patches) == 21 # one bar per lag 0..20 + + ax2 = ebm.viz.trace_plot(good, dim=1, max_chains=4) + assert len(ax2.lines) == 4 # capped at max_chains + + # Rank uniformity: matched chains give ~flat per-chain rank histograms; an + # offset chain concentrates in the top ranks. + ranks = ebm.viz._chain_ranks(good, dim=0) + m, n = ranks.shape + counts = torch.histc(ranks[0], bins=10, min=1, max=m * n) + assert counts.std().item() < 0.5 * (n / 10) # roughly uniform + + offset = torch.cat([torch.randn(5, 500, 2), 5.0 + torch.randn(1, 500, 2)], dim=0) + off_ranks = ebm.viz._chain_ranks(offset, dim=0) + top = (off_ranks[-1] > 0.9 * off_ranks.numel()).float().mean().item() + assert top > 0.5 # the shifted chain owns the top ranks + ax3 = ebm.viz.rank_plot(offset, dim=0, bins=10) + assert len(ax3.patches) > 0