Handle invalid/misaligned Chain labels in df_with_chain2array()

Author: jgabry

NEWS.md

@@ -1,6 +1,8 @@
 # bayesplot (development version)
 
-* Validate equal chain lengths in `validate_df_with_chain()`.
+* Validate equal chain lengths in `validate_df_with_chain()`, reject missing
+  chain labels, and renumber data-frame chain labels internally when converting
+  to arrays.
 * Added unit tests for previously untested edge cases in `param_range()`, `param_glue()`, and `tidyselect_parameters()` (no-match, partial-match, and negation behavior).
 * Bumped minimum version for `rstantools` from `>= 1.5.0` to `>= 2.0.0` .
 * Use `rlang::warn()` and `rlang::inform()` for selected PPC user messages instead of base `warning()` and `message()`.

R/helpers-mcmc.R

@@ -210,6 +210,9 @@ validate_df_with_chain <- function(x) {
     x$chain <- NULL
   }
   x$Chain <- as.integer(x$Chain)
+  if (anyNA(x$Chain)) {
+    abort("Chain values must not be NA.")
+  }
   rows_per_chain <- table(x$Chain)
   if (length(unique(rows_per_chain)) != 1) {
     abort("All chains must have the same number of iterations.")
@@ -222,6 +225,8 @@ validate_df_with_chain <- function(x) {
 df_with_chain2array <- function(x) {
   x <- validate_df_with_chain(x)
   chain <- x$Chain
+  # Renumber arbitrary chain labels to the contiguous 1:N indices used internally.
+  chain <- match(chain, sort(unique(chain)))
   n_chain <- length(unique(chain))
   a <- x[, !colnames(x) %in% "Chain", drop = FALSE]
   parnames <- colnames(a)

R/mcmc-overview.R

@@ -24,7 +24,9 @@
 #' frame with one column per parameter (if only a single chain or all chains
 #' have already been merged), or a data frame with one column per parameter plus
 #' an additional column `"Chain"` that contains the chain number (an integer)
-#' corresponding to each row in the data frame.
+#' corresponding to each row in the data frame. When a `"Chain"` column is
+#' supplied, each chain must have the same number of iterations. Chain labels
+#' are used to identify groups and are renumbered internally to `1:N`.
 #' * __draws__: Any of the `draws` formats supported by the
 #' \pkg{posterior} package.
 #'

man/MCMC-overview.Rd

@@ -106,6 +106,14 @@ test_that("validate_df_with_chain works", {
   tbl <- tibble::tibble(parameter=rnorm(n=40), Chain=rep(1:4, each=10))
   a <- validate_df_with_chain(tbl)
   expect_type(a$Chain, "integer")
+
+  missing_chain_df <- data.frame(
+    Chain = c(1L, 1L, NA_integer_, NA_integer_),
+    V1 = rnorm(4),
+    V2 = rnorm(4)
+  )
+  expect_error(validate_df_with_chain(missing_chain_df),
+               "Chain values must not be NA")
 })
 
 test_that("df_with_chain2array works", {
@@ -124,6 +132,17 @@ test_that("df_with_chain2array works", {
                "All chains must have the same number of iterations")
   expect_error(df_with_chain2array(unequal_df),
                "All chains must have the same number of iterations")
+
+  renumbered_df <- data.frame(
+    Chain = c(2L, 2L, 3L, 3L),
+    V1 = 1:4,
+    V2 = 5:8
+  )
+  a <- df_with_chain2array(renumbered_df)
+  expect_equal(dim(a), c(2, 2, 2))
+  expect_identical(unname(a[, 1, "V1"]), c(1L, 2L))
+  expect_identical(unname(a[, 2, "V1"]), c(3L, 4L))
+  expect_identical(as.character(dimnames(a)$Chain), c("1", "2"))
 })
 
 
@@ -316,6 +335,7 @@ test_that("diagnostic_factor.rhat works", {
   )
   expect_identical(levels(r), c("low", "ok", "high"))
 })
+
 test_that("diagnostic_factor.neff_ratio works", {
   ratios <- new_neff_ratio(c(low = 0.05, low = 0.01,
                              ok = 0.2, ok = 0.49,