Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

reference updates #374

Merged
merged 3 commits into from
Aug 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 11 additions & 11 deletions R/nested_rhat.R
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
#' Nested Rhat convergence diagnostic
#'
#' Compute the nested Rhat convergence diagnostic for a single
#' variable as proposed in Margossian et al. (2023).
#' variable as proposed in Margossian et al. (2024).
#'
#' @family diagnostics
#' @template args-conv
Expand All @@ -20,11 +20,11 @@
#' Note that there is a slight difference in the calculation of Rhat
#' and nested Rhat, as nested Rhat is lower bounded by 1. This means
#' that nested Rhat with one chain per superchain will not be
#' exactly equal to basic Rhat (see Footnote 1 in Margossian et
#' al. (2023)).
#' exactly equal to basic Rhat (see Footnote 3 in Margossian et
#' al. (2024)).
#'
#' @template return-conv
#' @template ref-margossian-nestedrhat-2023
#' @template ref-margossian-nestedrhat-2024
#'
#' @examples
#' mu <- extract_variable_matrix(example_draws(), "mu")
Expand Down Expand Up @@ -78,17 +78,17 @@ rhat_nested.rvar <- function(x, superchain_ids, ...) {

superchains <- unique(superchain_ids)

# mean and variance of chains calculated as in rhat
# mean and variance of chains calculated as in Rhat
chain_mean <- matrixStats::colMeans2(x)
chain_var <- matrixStats::colVars(x, center = chain_mean)

# mean of superchains calculated by only including specified chains
# (equation 15 in Margossian et al. 2023)
# (equation 4 in Margossian et al. 2024)
superchain_mean <- sapply(
superchains, function(k) mean(x[, which(superchain_ids == k)])
)

# between-chain variance estimate (B_k in equation 18 in Margossian et al. 2023)
# between-chain variance estimate (Bhat_k in equation 7 in Margossian et al. 2024)
if (nchains_per_superchain == 1) {
var_between_chain <- 0
} else {
Expand All @@ -98,7 +98,7 @@ rhat_nested.rvar <- function(x, superchain_ids, ...) {
)
}

# within-chain variance estimate (W_k in equation 18 in Margossian et al. 2023)
# within-chain variance estimate (What_k in equation 7 in Margossian et al. 2024)
if (niterations == 1) {
var_within_chain <- 0
} else {
Expand All @@ -108,12 +108,12 @@ rhat_nested.rvar <- function(x, superchain_ids, ...) {
)
}

# between-superchain variance (nB in equation 17 in Margossian et al. 2023)
# between-superchain variance (Bhat_nu in equation 6 in Margossian et al. 2024)
var_between_superchain <- var(superchain_mean)

# within-superchain variance (nW in equation 18 in Margossian et al. 2023)
# within-superchain variance (What_nu in equation 7 in Margossian et al. 2024)
var_within_superchain <- mean(var_within_chain + var_between_chain)

# nested Rhat (nRhat in equation 19 in Margossian et al. 2023)
# nested Rhat (Rhat_nu in equation 8 in Margossian et al. 2024)
sqrt(1 + var_between_superchain / var_within_superchain)
}
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
#' @references
#' Charles C. Margossian, Matthew D. Hoffman, Pavel Sountsov, Lionel
#' Riou-Durand, Aki Vehtari and Andrew Gelman (2023). Nested R-hat:
#' Riou-Durand, Aki Vehtari and Andrew Gelman (2024). Nested R-hat:
#' Assessing the convergence of Markov chain Monte Carlo when running
#' many short chains. arxiv:arXiv:2110.13017 (version 4)
#' many short chains. *Bayesian Analysis*. doi:10.1214/24-BA1453
2 changes: 1 addition & 1 deletion man-roxygen/ref-vehtari-rhat-2021.R
Original file line number Diff line number Diff line change
Expand Up @@ -2,5 +2,5 @@
#' Aki Vehtari, Andrew Gelman, Daniel Simpson, Bob Carpenter, and
#' Paul-Christian Bürkner (2021). Rank-normalization, folding, and
#' localization: An improved R-hat for assessing convergence of
#' MCMC (with discussion). *Bayesian Data Analysis*. 16(2), 667-–718.
#' MCMC (with discussion). *Bayesian Analysis*. 16(2), 667-–718.
#' doi:10.1214/20-BA1221
2 changes: 1 addition & 1 deletion man/ess_basic.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion man/ess_bulk.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion man/ess_quantile.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion man/ess_sd.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion man/ess_tail.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion man/mcse_quantile.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion man/mcse_sd.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion man/rhat.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion man/rhat_basic.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

10 changes: 5 additions & 5 deletions man/rhat_nested.Rd

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading