Fix df_with_chain2array() silently recycling data with unequal chain lengths

Author: utkarshpawade

NEWS.md

@@ -1,5 +1,6 @@
 # bayesplot (development version)
 
+* Fix `df_with_chain2array()` silently recycling data when chains have unequal iterations.
 * 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

@@ -222,7 +222,12 @@ df_with_chain2array <- function(x) {
   a <- x[, !colnames(x) %in% "Chain", drop = FALSE]
   parnames <- colnames(a)
   a <- as.matrix(a)
-  x <- array(NA, dim = c(ceiling(nrow(a) / n_chain), n_chain, ncol(a)))
+  rows_per_chain <- table(chain)
+  if (length(unique(rows_per_chain)) != 1) {
+    abort("All chains must have the same number of iterations.")
+  }
+  n_iter <- as.integer(rows_per_chain[[1]])
+  x <- array(NA, dim = c(n_iter, n_chain, ncol(a)))
   for (j in seq_len(n_chain)) {
     x[, j, ] <- a[chain == j,, drop=FALSE]
   }

tests/testthat/test-helpers-mcmc.R

@@ -113,6 +113,15 @@ test_that("df_with_chain2array works", {
   expect_mcmc_array(a)
 
   expect_error(df_with_chain2array(dframe), "is_df_with_chain")
+
+  # Unequal chain lengths should error, not silently recycle
+  unequal_df <- data.frame(
+    Chain = c(1L, 1L, 1L, 1L, 2L, 2L, 2L),
+    V1 = rnorm(7),
+    V2 = rnorm(7)
+  )
+  expect_error(df_with_chain2array(unequal_df),
+               "All chains must have the same number of iterations")
 })