arviz_stats.diagnose

Contents

arviz_stats.diagnose#

arviz_stats.diagnose(data, *, var_names=None, filter_vars=None, coords=None, sample_dims=None, group='posterior', rhat_max=1.01, ess_min_ratio=0.001, bfmi_threshold=0.3, show_diagnostics=True, return_diagnostics=False)[source]#

Run comprehensive diagnostic checks for MCMC sampling.

This function performs diagnostic checks on MCMC samples similar to CmdStan’s diagnose utility. It checks for:

  • Divergent transitions

  • Maximum tree depth saturation

  • Low E-BFMI (Energy Bayesian Fraction of Missing Information)

  • Low effective sample size (ESS)

  • High R-hat values

See [1] and [2] for more details. You can also check https://arviz-devs.github.io/EABM/Chapters/MCMC_diagnostics.html for a more practical overview.

Parameters:
dataxarray.DataTree, xarray.Dataset, or InferenceData-like

Input data. To be able to compute all diagnostics, the data should contain MCMC posterior samples (see group argument) and sampler statistics (in “sample_stats” group).

var_namesstr or list of str, optional

Names of variables to check for R-hat and ESS diagnostics. If None, checks all variables.

filter_vars{None, “like”, “regex”}, default None

How to filter variable names. See filter_vars for details.

coordsdict, optional

Coordinates to select a subset of the data.

sample_dimsiterable of hashable, optional

Dimensions to be considered sample dimensions. Default from rcParams["data.sample_dims"].

groupstr, default “posterior”

Group to check for convergence diagnostics (R-hat, ESS).

rhat_maxfloat, default 1.01

Maximum acceptable R-hat value. Parameters with R-hat > rhat_max will be flagged.

ess_min_ratiofloat, default 0.001

Minimum acceptable ratio of ESS to total samples. Parameters with ESS/N < ess_min_ratio will be flagged. A flag is also emitted if ESS is lower than 100 * number of chains.

bfmi_thresholdfloat, default 0.3

Minimum acceptable E-BFMI value. Values below this threshold indicate potential issues with the sampler’s exploration.

show_diagnosticsbool, default True

If True, print diagnostic messages to stdout. If False, return results silently.

return_diagnosticsbool, default False

If True, return a dictionary with detailed diagnostic results in addition to the boolean has_errors flag.

Returns:
has_errorsbool

True if any diagnostic checks failed, False otherwise.

diagnosticsdict, optional

Only returned if return_diagnostics=True.

  • “divergent”: dict with keys “n_divergent”, “pct”, “total_samples”

  • “treedepth”: dict with keys “n_max”, “pct”, “total_samples”

  • “bfmi”: dict with keys “bfmi_values”, “failed_chains”, “threshold”

  • “ess”: dict with keys “bad_params”, “ess_values”, “threshold_ratio”, “total_samples”

  • “rhat”: dict with keys “bad_params”, “rhat_values”, “threshold”

See also

rhat

Compute R-hat convergence diagnostic

ess

Compute effective sample size

bfmi

Compute Bayesian fraction of missing information

summary

Create a data frame with summary statistics, including diagnostics.

References

[1]

Vehtari et al. Rank-normalization, folding, and localization: An improved Rhat for assessing convergence of MCMC. Bayesian Analysis. 16(2) (2021) https://doi.org/10.1214/20-BA1221. arXiv preprint https://arxiv.org/abs/1903.08008

[2]

Betancourt. Diagnosing Suboptimal Cotangent Disintegrations in Hamiltonian Monte Carlo. (2016) https://arxiv.org/abs/1604.00695

Examples

Get diagnostics printted to stdout:

In [1]: import arviz_stats as azs
   ...: from arviz_base import load_arviz_data
   ...: data = load_arviz_data('centered_eight')
   ...: azs.diagnose(data)
   ...: 
Divergences
19 of 2000 (0.95%) transitions ended with a divergence.
These divergent transitions indicate that HMC is not fully able to explore the posterior distribution.
Try increasing adapt delta closer to 1.
If this doesn't remove all divergences, try to reparameterize the model.

Tree depth
Treedepth satisfactory for all transitions.

E-BFMI
Chain 0: E-BFMI = 0.287
E-BFMI values are below the threshold 0.30 which suggests that HMC may have trouble exploring the target distribution.
If possible, try to reparameterize the model.

ESS
The following parameters has fewer than 400 effective samples:
  mu, theta, tau
This suggests that the sampler may not have fully explored the posterior distribution for this parameter.
Consider reparameterizing the model or increasing the number of samples.

R-hat
The following parameters has R-hat values greater than 1.01:
  mu, theta, tau
Such high values indicate incomplete mixing and biased estimation.
You should consider regularizing your model with additional prior information or a more effective parameterization.
Out[1]: True

Get detailed diagnostic information without printing messages:

In [2]: _, diagnostics = azs.diagnose(data, return_diagnostics=True, show_diagnostics=False)
   ...: diagnostics
   ...: 
Out[2]: 
{'divergent': {'n_divergent': 19,
  'pct': np.float64(0.95),
  'total_samples': np.int64(2000)},
 'treedepth': {'n_max': 0,
  'pct': np.float64(0.0),
  'total_samples': np.int64(2000)},
 'bfmi': {'bfmi_values': array([0.28671974, 0.37830054, 0.34555751, 0.32426828]),
  'failed_chains': [0],
  'threshold': 0.3},
 'ess': {'bad_params': [],
  'ess_values': <xarray.Dataset> Size: 0B
  Dimensions:  ()
  Data variables:
      *empty*,
  'threshold_ratio': 0.001,
  'total_samples': np.int64(2000)},
 'rhat': {'bad_params': ['mu', 'theta', 'tau'],
  'rhat_values': <xarray.Dataset> Size: 592B
  Dimensions:  (school: 8)
  Coordinates:
    * school   (school) <U16 512B 'Choate' 'Deerfield' ... 'Mt. Hermon'
  Data variables:
      mu       float64 8B 1.025
      theta    (school) float64 64B 1.007 1.011 1.01 1.01 1.019 1.012 1.012 1.012
      tau      float64 8B 1.028,
  'threshold': 1.01}}
Contents