mixle.inference.diagnostics module

Public MCMC convergence diagnostics over plain arrays.

These functions are deliberately free of any MCMCResult / _Slot dependency: they operate on ordinary NumPy arrays so the draws of any sampler (mixle’s or an external one) can be diagnosed. They are the generic cores behind mixle.inference.mcmc.gelman_rubin and MCMCResult.effective_sample_size.

rhat(chains)[source]

Gelman-Rubin potential scale reduction factor (R-hat) per parameter dimension.

R-hat compares the variance between chains to the variance within chains. Values near 1.0 indicate the chains have mixed and are sampling a common target; values above ~1.01-1.1 flag non-convergence. Standard multi-chain check (Gelman & Rubin 1992).

Parameters:

chains (Any) – array shaped (n_chains, n_draws) or (n_chains, n_draws, d). Needs at least two chains and two draws.

Returns:

A length-d array of R-hat values (one per parameter).

Return type:

ndarray

ess(samples, max_lag=None)[source]

Effective sample size per parameter from positive autocorrelation lags.

Accepts a single chain (n_draws,) / (n_draws, d) or a stack of chains (n_chains, n_draws, d) (the chains are pooled into one autocorrelation estimate after centering each chain by its own mean). Uses the initial-positive-sequence truncation (Geyer): sum autocorrelations until the first non-positive lag.

Returns:

A length-d array of ESS values.

Parameters:
  • samples (Any)

  • max_lag (int | None)

Return type:

ndarray

split_rhat(chains)[source]

Rank-normalized split-R-hat (Vehtari et al. 2021) – the robust, recommended R-hat.

Splits each chain in half (doubling the chain count, so within-chain non-stationarity is caught), rank-normalizes the pooled draws (robust to heavy tails / non-normality), then applies the Gelman-Rubin R-hat. Convergence is typically declared at < 1.01.

Parameters:

chains (Any)

Return type:

ndarray

ess_bulk(chains)[source]

Bulk effective sample size: ESS of the rank-normalized split chains (mixing in the distribution body).

Parameters:

chains (Any)

Return type:

ndarray

ess_tail(chains, prob=0.05)[source]

Tail effective sample size: the smaller ESS of the lower-/upper-prob tail indicators.

Tail quantiles mix more slowly than the bulk; tail-ESS is min(ESS[1(x <= q_prob)], ESS[1(x >= q_{1-prob})]) over the split chains (Vehtari et al. 2021), surfacing poor tail exploration that bulk-ESS misses.

Parameters:
Return type:

ndarray

folded_split_rhat(chains)[source]

Folded rank-normalized split-R-hat (Vehtari et al. 2021): split-R-hat on |x - median(x)|.

The plain split_rhat() compares chain locations; folding about the median makes it compare chain scales/tails, catching the case where chains share a mean but differ in spread. Use together with split_rhat() (or rhat_max()).

Parameters:

chains (Any)

Return type:

ndarray

rhat_max(chains)[source]

The recommended convergence R-hat: max(split_rhat, folded_split_rhat) (Vehtari et al. 2021).

A single number that flags non-convergence in either the location (bulk) or the scale (folded) of the chains; declare convergence at < 1.01.

Parameters:

chains (Any)

Return type:

ndarray

mcse_mean(chains)[source]

Monte Carlo standard error of the posterior mean: sd(x) / sqrt(ESS) per parameter.

The sampling error in the estimated mean from autocorrelated draws – the posterior standard deviation deflated by the (autocorrelation-based) effective sample size of the raw chains.

Parameters:

chains (Any)

Return type:

ndarray

mcmc_summary(chains)[source]

Per-parameter posterior + convergence summary (the ArviZ-style summary table).

Returns one dict per parameter with the posterior mean/sd and 5/50/95% quantiles, plus the recommended convergence diagnostics: r_hat (= rhat_max(), the max of the bulk and folded rank-normalized split-R-hats), ess_bulk, ess_tail, and mcse_mean (Vehtari et al. 2021). Declare convergence at r_hat < 1.01 with ess_bulk/ess_tail comfortably in the hundreds.

Parameters:

chains (Any)

Return type:

list[dict[str, float]]

geweke_z(chain, first=0.1, last=0.5)[source]

Geweke convergence z-score comparing the start and end of a single chain (per parameter).

Tests stationarity within one chain (Geweke 1992): if the chain has converged, the mean of its first first fraction and its last last fraction agree, so the z-statistic (mean_a - mean_b) / sqrt(se_a^2 + se_b^2) – with autocorrelation-adjusted standard errors (the segment variance divided by its effective sample size) – is approximately standard normal. |z| < 2 indicates convergence; a large |z| flags a chain still drifting. Returns one z per parameter. Complements the multi-chain split_rhat().

Parameters:
Return type:

ndarray