Sec. VII: Stochastic Convergence Diagnostics (jaxpe.diagnostics)

1. Split R-hat Convergence Criterion
2. Autocovariance and Effective Sample Size
3. High-Dimensional Visualization
4. API Reference
   1. split_rhat
   2. effective_sample_size
   3. js_divergence
   4. corner_plot
   5. trace_plot
   6. REFERENCES

This section details the statistical validation of the generated Markov chain ensembles. A Markov chain is defined as a sequence of correlated samples; distinguishing transient, non-stationary adaptation phases from ergodic mixing across the invariant measure \(\pi(x)\) requires robust convergence diagnostics.

Split R-hat Convergence Criterion

The Split R-hat (\(\hat{R}\)) statistic [1, 2] evaluates the geometric convergence of the ensemble by rigorously comparing the variance between \(M\) independent Markov chains against the variance within the individual chains.

Let \(\theta^\mu_{(m, i)}\) denote the \(i\)-th sample of the \(\mu\)-th parameter in the \(m\)-th chain, where \(i = 1 \dots N\). To detect non-stationarity within a single chain, we split each sequence in half, yielding \(2M\) chains of length \(N/2\).

The empirical between-chain variance scalar \(B\) and within-chain variance scalar \(W\) for a parameter \(\theta\) are defined explicitly as:

\[B = \frac{N/2}{2M - 1} \sum_{m=1}^{2M} \left( \frac{2}{N} \sum_{i=1}^{N/2} \theta_{(m,i)} - \frac{1}{MN} \sum_{m=1}^{2M} \sum_{i=1}^{N/2} \theta_{(m,i)} \right)^2\] \[W = \frac{1}{2M (N/2 - 1)} \sum_{m=1}^{2M} \sum_{i=1}^{N/2} \left( \theta_{(m,i)} - \frac{2}{N} \sum_{j=1}^{N/2} \theta_{(m,j)} \right)^2\]

Under the assumption of target ergodicity, the true marginal posterior variance \(\Sigma = \text{Var}(\theta)\) is estimated as an unbiased convex combination of these empirical moments:

\[\widehat{\Sigma} = \frac{N/2 - 1}{N/2} W + \frac{1}{N/2} B\]

The potential scale reduction factor evaluates the ratio of this pooled variance to the within-chain variance:

\[\hat{R} = \sqrt{\frac{\widehat{\Sigma}}{W}}\]

As \(N \to \infty\), asymptotic stationarity dictates that \(B \to W\), driving \(\hat{R} \to 1\). In practical gravitational-wave inference, any value \(\hat{R} > 1.05\) indicates a severe failure of global mixing, signaling that the disparate chains occupy distinct, disconnected topological modes.

In jaxpe, this calculation is exposed via split_rhat:

from jaxpe.diagnostics.stats import split_rhat

# chains shape: (n_chains, n_samples, n_dim)
rhat_values = split_rhat(chains)

Autocovariance and Effective Sample Size

Because consecutive states generated by continuous Langevin or Hamiltonian flows are strictly correlated, the true statistical power of an \(N\)-length Markov chain is sharply reduced. We quantify this by analyzing the autocovariance function of the stochastic process.

For a stationary chain, the autocovariance at lag \(t\) is \(\gamma_t = \text{Cov}(\theta_{(i)}, \theta_{(i+t)})\). The normalized autocorrelation coefficient is \(\rho_t = \gamma_t / \gamma_0\).

The variance of the empirical mean \(\bar{\theta}\) scales not by \(1/N\), but by the infinite sum of these correlations. We define the Effective Sample Size (ESS), \(N_{\text{eff}}\), as the equivalent number of strictly independent samples drawn from the invariant measure:

\[N_{\text{eff}} = \frac{MN}{1 + 2 \sum_{t=1}^{\infty} \hat{\rho}_t}\]

where \(\hat{\rho}_t\) is the computationally truncated estimate of the autocorrelation. A robust \(N_{\text{eff}}\) is the primary prerequisite for asserting that the Monte Carlo standard errors \(\sigma_{\text{MC}} = \sqrt{\widehat{\Sigma} / N_{\text{eff}}}\) are sufficiently bounded to constrain astrophysics.

This critical metric is computed via effective_sample_size:

from jaxpe.diagnostics.stats import effective_sample_size

# chains shape: (n_chains, n_samples, n_dim)
ess_values = effective_sample_size(chains)

High-Dimensional Visualization

The diagnostics module leverages corner.py to project the high-dimensional measure \(\pi(\theta^\mu \mid d)\) onto 1D and 2D marginal distributions. These visual projections remain the most intuitive mechanism for exposing the complex Riemann curvature (e.g., severe mass-spin degeneracies) of the physical parameter manifold.

API Reference

split_rhat

jaxpe.diagnostics.stats.split_rhat(xs) Computes the potential scale reduction factor (Split R-hat) given an array of samples shaped (n_chains, n_samples, n_dim). Values significantly greater than 1.0 indicate non-stationarity.

effective_sample_size

jaxpe.diagnostics.stats.effective_sample_size(xs) Computes the Effective Sample Size (ESS) for an array of MCMC chains by summing over the integrated autocorrelation time.

js_divergence

jaxpe.diagnostics.metrics.js_divergence(a, b, bins=...)

Jensen–Shannon divergence (base 2) between two 1-D sample sets — the workhorse for asking whether two samplers, or two runs of one sampler, found the same posterior:

\[{\rm JS}(P \parallel Q) = \tfrac12 D_{\rm KL}(P \parallel M) + \tfrac12 D_{\rm KL}(Q \parallel M), \qquad M = \tfrac12 (P + Q).\]

Unlike the KL divergence it is symmetric, bounded (in \([0, 1]\) base 2), and finite even when the supports differ — all three matter for finite MCMC samples.

Interpret it only against a null. Any two finite samples of the same distribution have \({\rm JS} > 0\), so a raw value is close to meaningless. The honest reference is a floor built by splitting one sampler’s own chains in half and computing the same statistic; the comparison worth reporting is the ratio to that floor. A JS that looks alarming at \(5\times10^{-3}\) may sit at or below the Monte-Carlo floor for the sample size in hand — this is worked through concretely in the BNS benchmark, where it reversed a published conclusion. The estimator is histogram-based, so bins is a real analysis choice, not a formality.

corner_plot

jaxpe.diagnostics.plots.corner_plot(samples, names=..., truths=...) Pair grid of 1-D and 2-D marginal posteriors, with optional truth markers.

trace_plot

jaxpe.diagnostics.plots.trace_plot(xs, names=..., max_chains=...) Chain traces against step index — the fastest way to see stuck chains, a transient that has not decayed, or multimodality that \(\widehat{R}\) alone may miss. max_chains caps how many are drawn, since overplotting 128 chains hides exactly what the plot is for.

A caution these three share. \(\widehat{R}\) measures whether chains agree with each other, which is not the same as agreeing with the target. A biased-but-consistent sampler — an unadjusted integrator at fixed step size, say — can pass any \(\widehat{R}\)/ESS gate while converging to the wrong distribution. Passing the gate is necessary, not sufficient; js_divergence against a trusted reference is what tests the second thing.

---

REFERENCES

[1] A. Gelman and D. B. Rubin, “Inference from Iterative Simulation Using Multiple Sequences,” Stat. Sci. 7, 457 (1992).

[2] A. Vehtari, A. Gelman, D. Simpson, B. Carpenter, and P. C. Bürkner, “Rank-Normalization, Folding, and Localization: An Improved \(\widehat{R}\) for Assessing Convergence of MCMC,” Bayesian Anal. 16, 667 (2021).