Skip to content

Commit

Permalink
Merge pull request #374 from n-kall/master
Browse files Browse the repository at this point in the history
reference updates
  • Loading branch information
paul-buerkner authored Aug 6, 2024
2 parents 18c915e + 88b95de commit 20b6d3a
Show file tree
Hide file tree
Showing 13 changed files with 28 additions and 28 deletions.
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.

0 comments on commit 20b6d3a

Please sign in to comment.