blackjax.diagnostics

blackjax.diagnostics#

MCMC diagnostics.

Functions#

potential_scale_reduction(→ blackjax.types.Array)

Gelman and Rubin (1992)'s potential scale reduction for computing multiple MCMC chain convergence.

effective_sample_size(→ blackjax.types.Array)

Compute estimate of the effective sample size (ess).

Module Contents#

potential_scale_reduction(input_array: blackjax.types.ArrayLike, chain_axis: int = 0, sample_axis: int = 1) → blackjax.types.Array[source]#

Gelman and Rubin (1992)’s potential scale reduction for computing multiple MCMC chain convergence.

Parameters:

  • input_array – An array representing multiple chains of MCMC samples. The array must contains a chain dimension and a sample dimension.

  • chain_axis – The axis indicating the multiple chains. Default to 0.

  • sample_axis – The axis indicating a single chain of MCMC samples. Default to 1.

Return type:

NDArray of the resulting statistics (r-hat), with the chain and sample dimensions squeezed.

Notes

The diagnostic is computed by:

\[\hat{R} = \frac{\hat{V}}{W}\]

where \(W\) is the within-chain variance and \(\hat{V}\) is the posterior variance estimate for the pooled traces. This is the potential scale reduction factor, which converges to unity when each of the traces is a sample from the target posterior. Values greater than one indicate that one or more chains have not yet converged [stac, GR92].

effective_sample_size(input_array: blackjax.types.ArrayLike, chain_axis: int = 0, sample_axis: int = 1) → blackjax.types.Array[source]#

Compute estimate of the effective sample size (ess).

Parameters:

  • input_array – An array representing multiple chains of MCMC samples. The array must contains a chain dimension and a sample dimension.

  • chain_axis – The axis indicating the multiple chains. Default to 0.

  • sample_axis – The axis indicating a single chain of MCMC samples. Default to 1.

Return type:

NDArray of the resulting statistics (ess), with the chain and sample dimensions squeezed.

Notes

The basic ess (\(N_{\mathit{eff}}\)) diagnostic is computed by:

\[\hat{N}_{\mathit{eff}} = \frac{MN}{\hat{\tau}}\]
\[\hat{\tau} = -1 + 2 \sum_{t'=0}^K \hat{P}_{t'}\]

where \(M\) is the number of chains, \(N\) the number of draws, \(\hat{\rho}_t\) is the estimated _autocorrelation at lag \(t\), and \(K\) is the last integer for which \(\hat{P}_{K} = \hat{\rho}_{2K} + \hat{\rho}_{2K+1}\) is still positive [staa, GCSR95].

The current implementation is similar to Stan, which uses Geyer’s initial monotone sequence criterion [Gey92, Gey11].