Allow NAs in prepare_mcmc_array

NEWS.md

@@ -1,5 +1,6 @@
 # bayesplot (development version)
 
+* `prepare_mcmc_array()` now warns instead of erroring on `NA`s in the input.
 * Fixed `validate_chain_list()` colnames check to compare all chains, not just the first two.
 * Added test verifying `legend_move("none")` behaves equivalently to `legend_none()`.
 * Added singleton-dimension edge-case tests for exported `_data()` functions.

R/helpers-mcmc.R

@@ -29,7 +29,9 @@ prepare_mcmc_array <- function(x,
     abort("Arrays should have 2 or 3 dimensions. See help('MCMC-overview').")
   }
   if (anyNA(x)) {
-    abort("NAs not allowed in 'x'.")
+    warn(
+      "NAs found in 'x'. These are passed through as-is and may affect the resulting plots."
+    )
   }
 
   if (rlang::is_quosures(pars)) {

tests/testthat/test-helpers-mcmc.R

@@ -252,9 +252,12 @@ test_that("transformations recycled properly if not a named list", {
 
 
 # prepare_mcmc_array ------------------------------------------------------
-test_that("prepare_mcmc_array errors if NAs", {
-  arr[1,1,1] <- NA
-  expect_error(prepare_mcmc_array(arr), "NAs not allowed")
+test_that("prepare_mcmc_array warns but does not error if NAs", {
+  arr_na <- arr
+  arr_na[1, 1, 1] <- NA
+  expect_warning(out <- prepare_mcmc_array(arr_na), "NAs found in 'x'")
+  expect_s3_class(out, "mcmc_array")
+  expect_true(anyNA(out))
 })
 test_that("prepare_mcmc_array processes non-array input types correctly", {
   # errors are mostly covered by tests of the many internal functions above

tests/testthat/test-mcmc-traces.R

@@ -165,6 +165,33 @@ test_that("mcmc_trace renders correctly", {
   vdiffr::expect_doppelganger("mcmc_trace (iter1 offset)", p_iter1)
 })
 
+# https://github.com/stan-dev/bayesplot/issues/250
+test_that("mcmc_trace renders correctly with NAs in draws", {
+  testthat::skip_on_cran()
+  testthat::skip_if_not_installed("vdiffr")
+  skip_on_r_oldrel()
+
+  set.seed(250)
+  draws_full_na <- array(
+    rnorm(500 * 4 * 2),
+    dim = c(500, 4, 2),
+    dimnames = list(NULL, NULL, c("theta[1,3]", "theta[2,3]"))
+  )
+  draws_full_na[, , "theta[2,3]"] <- NA
+
+  draws_partial_na <- draws_full_na
+  draws_partial_na[, , "theta[2,3]"] <- rnorm(500 * 4)
+  draws_partial_na[10:100, , "theta[2,3]"] <- NA
+
+  suppressWarnings({
+    p_full_na <- mcmc_trace(draws_full_na)
+    p_partial_na <- mcmc_trace(draws_partial_na)
+  })
+
+  vdiffr::expect_doppelganger("mcmc_trace (NA parameter)", p_full_na)
+  vdiffr::expect_doppelganger("mcmc_trace (partial NA parameter)", p_partial_na)
+})
+
 test_that("mcmc_rank_overlay renders correctly", {
   testthat::skip_on_cran()
   testthat::skip_if_not_installed("vdiffr")