Move unequal chain length check to validate_df_with_chain()

Author: utkarshpawade

NEWS.md

@@ -1,6 +1,6 @@
 # bayesplot (development version)
 
-* Fix `df_with_chain2array()` silently recycling data when chains have unequal iterations.
+* Validate equal chain lengths in `validate_df_with_chain()`.
 * 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,10 @@ validate_df_with_chain <- function(x) {
     x$chain <- NULL
   }
   x$Chain <- as.integer(x$Chain)
+  rows_per_chain <- table(x$Chain)
+  if (length(unique(rows_per_chain)) != 1) {
+    abort("All chains must have the same number of iterations.")
+  }
   x
 }
 
@@ -222,11 +226,7 @@ df_with_chain2array <- function(x) {
   a <- x[, !colnames(x) %in% "Chain", drop = FALSE]
   parnames <- colnames(a)
   a <- as.matrix(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]])
+  n_iter <- nrow(a) %/% n_chain
   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

@@ -114,12 +114,14 @@ test_that("df_with_chain2array works", {
 
   expect_error(df_with_chain2array(dframe), "is_df_with_chain")
 
-  # Unequal chain lengths should error, not silently recycle
+  # Unequal chain lengths should error via validate_df_with_chain
   unequal_df <- data.frame(
     Chain = c(1L, 1L, 1L, 1L, 2L, 2L, 2L),
     V1 = rnorm(7),
     V2 = rnorm(7)
   )
+  expect_error(validate_df_with_chain(unequal_df),
+               "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")
 })