arviz_stats.rhat_nested

Contents

arviz_stats.rhat_nested#

arviz_stats.rhat_nested(data, sample_dims=None, group='posterior', var_names=None, filter_vars=None, method='rank', coords=None, superchain_ids=None, chain_axis=0, draw_axis=1)[source]#

Compute nested R-hat.

Nested R-hat is a convergence diagnostic useful when running many short chains. It is calculated on superchains, which are groups of chains that have been initialized at the same point.

Note that there is a slight difference in the calculation of R-hat and nested R-hat, as nested R-hat is lower bounded by 1. This means that nested R-hat with one chain per superchain will not be exactly equal to basic R-hat see [1] for details.

Parameters:
dataarray_like, xarray.DataArray, xarray.Dataset, xarray.DataTree, DataArrayGroupBy, DatasetGroupBy, or idata-like

Input data. It will have different pre-processing applied to it depending on its type:

  • array-like: call array layer within arviz-stats.

  • xarray object: apply dimension aware function to all relevant subsets

  • others: passed to arviz_base.convert_to_dataset

At least 2 superchains are needed to compute this diagnostic. Each superchain must contain at least 2 chains, or each chain must have at least 2 draws, so that the within-superchain variance can be estimated. The split-based methods (“rank”, “split”, “folded”, “z_scale”) split each chain in two and therefore need at least 2 draws per chain. method="identity" works with a single draw per chain and is the estimator exactly as defined in Margossian et al. (2024) [1].

sample_dimsiterable of hashable, optional

Dimensions to be considered sample dimensions and are to be reduced. Default rcParams["data.sample_dims"].

grouphashable, default “posterior”

Group on which to compute the R-hat.

var_namesstr or list of str, optional

Names of the variables for which the Rhat should be computed.

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

Valid methods are:

  • “rank” # recommended by Vehtari et al. (2021) [2]

  • “split”

  • “folded”

  • “z_scale”

  • “identity”

"identity" matches the definition in [1] and the posterior R package; the other methods first apply the split/rank/fold transformations of [2].

coordsdict, optional

Dictionary of dimension/index names to coordinate values defining a subset of the data for which to perform the computation.

superchain_idslist

List of length chains specifying which superchain each chain belongs to. There should be equal numbers of chains in each superchain. All chains within the same superchain are assumed to have been initialized at the same point.

chain_axis, draw_axisint, optional

Integer indicators of the axis that correspond to the chain and the draw dimension. chain_axis can be None.

Returns:
xarray.DataArray or xarray.Dataset

R-hat values for the specified variables and groups.

See also

rhat

Compute split R-hat convergence diagnostic

ess

Estimate the effective sample size (ess).

plot_forest

Forest plot to compare HDI intervals from a number of distributions.

References

[1] (1,2,3)

Margossian et al. Nested R-hat: Assessing the convergence of Markov Chain Monte Carlo when running many short chains. Bayesian Analysis, (2024). https://doi.org/10.1214/24-BA1453

[2] (1,2)

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

Examples

Calculate nested R-hat with two superchains of two chains each:

In [1]: from arviz_base import load_arviz_data
   ...: import arviz_stats as azs
   ...: data = load_arviz_data('centered_eight')
   ...: azs.rhat_nested(data, superchain_ids=[0, 0, 1, 1])
   ...: 
Out[1]: 
<xarray.DataTree 'posterior'>
Group: /posterior
    Dimensions:  (school: 8)
    Coordinates:
      * school   (school) <U16 512B 'Choate' 'Deerfield' ... 'Mt. Hermon'
    Data variables:
        mu       float64 8B 1.0
        theta    (school) float64 64B 1.002 1.001 1.001 1.001 1.001 1.003 1.004 1.0
        tau      float64 8B 1.005

With a single draw per chain the split-based methods return nan; use method="identity":

In [2]: import numpy as np
   ...: rng = np.random.default_rng(0)
   ...: samples = rng.normal(size=(8, 1))
   ...: azs.rhat_nested(samples, superchain_ids=[0, 0, 0, 0, 1, 1, 1, 1], method="identity")
   ...: 
Out[2]: array(1.07186933)
Contents