Call as.array on input to prepare_mcmc_array

Author: jgabry

R/helpers-mcmc.R

@@ -12,14 +12,27 @@ prepare_mcmc_array <- function(x,
     x <- df_with_chain2array(x)
   } else if (is_chain_list(x)) {
     x <- chain_list2array(x)
+  } else if (is.data.frame(x)) {
+    # data frame without Chain column
+    x <- as.matrix(x)
+  } else if (!is.array(x)) {
+    x <- as.array(x)
+  }
+
+  stopifnot(is.matrix(x) || is.array(x))
+  if (is.array(x) && !(length(dim(x)) %in% c(2,3))) {
+    stop("Arrays should have 2 or 3 dimensions. See help('MCMC-overview').")
+  }
+  if (anyNA(x)) {
+    stop("NAs not allowed in 'x'.")
   }
-  x <- validate_mcmc_x(x)
 
   parnames <- parameter_names(x)
   pars <- select_parameters(
     explicit = pars,
     patterns = regex_pars,
-    complete = parnames)
+    complete = parnames
+  )
 
   # possibly recycle transformations (apply same to all pars)
   if (is.function(transformations) ||
@@ -157,16 +170,12 @@ df_with_chain2array <- function(x) {
 # @param x object to check
 # @return TRUE or FALSE
 is_chain_list <- function(x) {
-  !is.data.frame(x) && is.list(x)
+  check1 <- !is.data.frame(x) && is.list(x)
+  dims <- sapply(x, function(chain) length(dim(chain)))
+  isTRUE(all(dims == 2))
 }
 
 validate_chain_list <- function(x) {
-  stopifnot(is_chain_list(x))
-  dims <- sapply(x, function(chain) length(dim(chain)))
-  if (!isTRUE(all(dims == 2))) {
-    stop("If 'x' is a list then all elements must be matrices.")
-  }
-
   n_chain <- length(x)
   for (i in seq_len(n_chain)) {
     nms <- colnames(as.matrix(x[[i]]))
@@ -276,25 +285,6 @@ STOP_need_multiple_chains <- function(call. = FALSE) {
 }
 
 
-# Perform some checks on user's 'x' input for MCMC plots
-#
-# @param x User's 'x' input to one of the mcmc_* functions.
-# @return x, unless an error is thrown.
-validate_mcmc_x <- function(x) {
-  stopifnot(!is_df_with_chain(x), !is_chain_list(x))
-  if (is.data.frame(x)) {
-    x <- as.matrix(x)
-  }
-
-  stopifnot(is.matrix(x) || is.array(x))
-  if (anyNA(x)) {
-    stop("NAs not allowed in 'x'.")
-  }
-
-  x
-}
-
-
 # Validate that transformations match parameter names
 validate_transformations <-
   function(transformations = list(),

R/mcmc-overview.R

@@ -18,9 +18,9 @@
 #'  corresponds to a Markov chain. All of the matrices should have the same
 #'  number of iterations (rows) and parameters (columns), and parameters should
 #'  have the same names and be in the same order.
-#'  \item \strong{matrix}: A \code{\link{matrix}} with one column per parameter.
-#'  If using matrix there should only be a single Markov chain or all chains
-#'  should already be merged (stacked).
+#'  \item \strong{matrix (2-D array)}: A \code{\link{matrix}} with one column
+#'  per parameter. If using matrix there should only be a single Markov chain or
+#'  all chains should already be merged (stacked).
 #'  \item \strong{data frame}: There are two types of \link[=data.frame]{data
 #'  frames} allowed. Either a data frame with one column per parameter (if only
 #'  a single chain or all chains have already been merged), or a data frame with

man/MCMC-overview.Rd

@@ -24,30 +24,6 @@ test_that("melt_mcmc does not convert integer parameter names to integers #162",
 })
 
 
-
-# validate_mcmc_x ----------------------------------------------------------
-test_that("validate_mcmc_x works", {
-  expect_identical(validate_mcmc_x(mat), mat)
-  expect_identical(validate_mcmc_x(mat1), mat1)
-  expect_identical(validate_mcmc_x(arr), arr)
-  expect_identical(validate_mcmc_x(arr1), arr1)
-  expect_identical(validate_mcmc_x(arr1chain), arr1chain)
-
-  # error if df_with_chain
-  expect_error(validate_mcmc_x(dframe_multiple_chains), "is_df_with_chain")
-
-  # converts regular df to matrix
-  expect_identical(validate_mcmc_x(dframe), as.matrix(dframe))
-
-  # NAs
-  mat[1, 2] <- NA
-  arr[1, 2, 3] <- NA
-  expect_error(validate_mcmc_x(mat), "NAs not allowed")
-  expect_error(validate_mcmc_x(arr), "NAs not allowed")
-})
-
-
-
 # 3-D array helpers --------------------------------------------------------
 test_that("is_mcmc_array works", {
   expect_false(is_mcmc_array(mat))
@@ -153,7 +129,6 @@ test_that("is_chain_list works", {
 })
 
 test_that("validate_chain_list works", {
-  expect_error(validate_chain_list(mat), "is_chain_list")
   expect_identical(validate_chain_list(chainlist), chainlist)
   expect_identical(validate_chain_list(chainlist1), chainlist1)
   expect_identical(validate_chain_list(chainlist1chain), chainlist1chain)
@@ -172,8 +147,6 @@ test_that("chain_list2array works", {
   expect_mcmc_array(chain_list2array(chainlist))
   expect_mcmc_array(chain_list2array(chainlist1))
   expect_mcmc_array(chain_list2array(chainlist1chain))
-
-  expect_error(chain_list2array(dframe), "is_chain_list")
 })
 
 
@@ -231,6 +204,42 @@ test_that("transformations recycled properly if not a named list", {
 })
 
 
+# prepare_mcmc_array ------------------------------------------------------
+test_that("prepare_mcmc_array processes non-array input types correctly", {
+  # errors are mostly covered by tests of the many internal functions above
+
+  # data frame with no Chain column (treat as 1 chain or merged chains)
+  a1 <- prepare_mcmc_array(dframe)
+  expect_s3_class(a1, "mcmc_array")
+  expect_equal(dim(a1), c(nrow(dframe), 1, ncol(dframe)))
+  expect_equal(parameter_names(a1), colnames(dframe))
+
+  # data frame with Chain column
+  a2 <- prepare_mcmc_array(dframe_multiple_chains)
+  expect_s3_class(a2, "mcmc_array")
+  n_chain <- max(dframe_multiple_chains$chain)
+  expect_equal(dim(a2), c(nrow(dframe) / n_chain, n_chain, ncol(dframe)))
+  expect_equal(parameter_names(a2), colnames(dframe))
+
+  # list of matrices with multiple chains
+  a3 <- prepare_mcmc_array(chainlist)
+  expect_s3_class(a3, "mcmc_array")
+  expect_equal(dim(a3), c(nrow(chainlist[[1]]), length(chainlist), ncol(chainlist[[1]])))
+  expect_equal(parameter_names(a3), colnames(chainlist[[1]]))
+
+  # object with acceptable as.array method
+  suppressPackageStartupMessages(library(rstanarm))
+  fit <- stan_glm(mpg ~ wt, data = mtcars, chains = 2, iter = 500, refresh = 0)
+  a4 <- prepare_mcmc_array(fit)
+  expect_s3_class(a4, "mcmc_array")
+  expect_equal(dim(a4), c(250, 2, 3))
+  expect_equal(parameter_names(a4), c("(Intercept)", "wt", "sigma"))
+
+  # object with unacceptable as.array method
+  fit2 <- lm(mpg ~ wt, data = mtcars)
+  expect_error(prepare_mcmc_array(fit2), "Arrays should have 2 or 3 dimensions.")
+})
+
 
 # rhat and neff helpers ---------------------------------------------------
 test_that("diagnostic_factor.rhat works", {