rhat_nested rejects valid superchain_ids that are not 0..K-1 (e.g. [1, 1, 2, 2])

[lucaluo925]

rhat_nested raises ValueError: Number of chains per superchain is not the same for each superchain for valid, equal-sized groupings whenever the labels are not exactly 0..K-1, e.g. 1-based ids [1, 1, 2, 2] (natural for R users), [0, 0, 2, 2] or [5, 5, 9, 9]. String labels fail inside np.bincount too.

Nested R-hat (Margossian et al. 2024, eq. 6-8) only depends on how chains are grouped, not on the label values, and the docstring only asks for a list of length chains with equal chains per superchain.

Cause: _rhat_nested in src/arviz_stats/base/diagnostics.py uses superchain_counts = np.bincount(superchain_ids) followed by a max == min check. bincount counts every integer from 0 to max(label), so unused labels get count 0. The rest of the function already loops over np.unique labels.

Suggested fix: _, superchain_counts = np.unique(superchain_ids, return_counts=True).

AI disclosure: I found this with the help of an AI assistant (Claude), which also ran the reproduction below and drafted this issue; I reviewed it before posting.

Reproduce

"""rhat_nested rejects valid superchain labels that are not 0..K-1.

Nested R-hat (Margossian et al. 2024, eq. 6-8) only depends on the *partition*
of chains into superchains, so relabelling must not change the result:
  Rhat_nu = sqrt(1 + B_nu / W_nu)
  B_nu = var_{ddof=1}(superchain means),  W_nu = mean_k( B_k + W_k )
"""
import numpy as np
import arviz_stats as azs
from arviz_base import from_dict
from arviz_stats.base import array_stats as A

rng = np.random.default_rng(1)
x = rng.normal(size=(4, 200))

def nested_rhat_ref(ary, ids):
    """Independent implementation (method='identity', no split / rank-normalization)."""
    ids = np.asarray(ids)
    labs = np.unique(ids)
    cm = ary.mean(axis=1); cv = ary.var(axis=1, ddof=1)
    sm = np.array([cm[ids == k].mean() for k in labs])
    Bk = np.array([cm[ids == k].var(ddof=1) for k in labs])
    Wk = np.array([cv[ids == k].mean() for k in labs])
    return np.sqrt(1 + sm.var(ddof=1) / np.mean(Bk + Wk))

for ids in ([0, 0, 1, 1], [1, 1, 0, 0], [1, 1, 2, 2], [0, 0, 2, 2], [5, 5, 9, 9]):
    ref = nested_rhat_ref(x, ids)
    try:
        got = A.rhat_nested(x, superchain_ids=ids, method="identity")
        print(ids, "arviz-stats:", round(float(got), 6), " reference:", round(ref, 6))
    except ValueError as e:
        print(ids, "arviz-stats: ValueError:", e, " reference:", round(ref, 6))

dt = from_dict({"posterior": {"mu": x}})
try:
    print("top-level, 1-based ids:", azs.rhat_nested(dt, superchain_ids=[1, 1, 2, 2]))
except ValueError as e:
    print("top-level azs.rhat_nested(ids=[1,1,2,2]): ValueError:", e)

Output:

[0, 0, 1, 1] arviz-stats: 1.000195  reference: 1.000195
[1, 1, 0, 0] arviz-stats: 1.000195  reference: 1.000195
[1, 1, 2, 2] arviz-stats: ValueError: Number of chains per superchain is not the same for each superchain  reference: 1.000195
[0, 0, 2, 2] arviz-stats: ValueError: Number of chains per superchain is not the same for each superchain  reference: 1.000195
[5, 5, 9, 9] arviz-stats: ValueError: Number of chains per superchain is not the same for each superchain  reference: 1.000195
top-level azs.rhat_nested(ids=[1,1,2,2]): ValueError: Number of chains per superchain is not the same for each superchain

Expected

All labellings of the same grouping give the same value (1.000195 here).

Environment

arviz-stats 1.4.0.dev0 (main, cefca73, 2026-09-24), arviz-base 1.3.0, numpy 2.5.3, Python 3.12.3.