Accept cmdstanvb, draws_array or draws_matrix as fitted_params

R/data.R

@@ -122,70 +122,106 @@ any_na_elements <- function(data) {
   any(has_na_elements)
 }
 
+#' Write posterior draws objects to csv files
+#' @noRd
+#' @param draws A `draws_array` from posterior pkg
+#' @param sampler_diagnostics Either `NULL` or a `draws_array` of sampler diagnostics
+#' @return Paths to CSV files (one per chain).
+#'
+draws_to_csv <- function(draws, sampler_diagnostics = NULL) {
+  n <- posterior::niterations(draws)
+  n_chains <- posterior::nchains(draws)
+  zeros <- rep(0, n * n_chains) # filler for creating dummy sampler diagnostics and lp__ if necessary
+  if (is.null(sampler_diagnostics)) {
+    # create dummy sampler diagnostics due to CmdStan requirement for all columns in GQ
+    sampler_diagnostics <- posterior::draws_array(
+      accept_stat__ = zeros,
+      stepsize__ = zeros,
+      treedepth__ = zeros,
+      n_leapfrog__ = zeros,
+      divergent__ = zeros,
+      energy__ = zeros,
+      .nchains = n_chains
+    )
+  }
+
+  # the columns must be in order "lp__, sampler_diagnostics, parameters"
+  draws_variables <- posterior::variables(draws)
+  if ("lp__" %in% draws_variables) {
+    lp__ <- NULL
+  } else { # create a dummy lp__ if it does not exist
+    lp__ <- posterior::draws_array(lp__ = zeros, .nchains = n_chains)
+  }
+  all_variables <- c("lp__", posterior::variables(sampler_diagnostics), draws_variables[!(draws_variables %in% c("lp__", "lp_approx__"))])
+  draws <- posterior::subset_draws(
+    posterior::bind_draws(draws, sampler_diagnostics, lp__, along = "variable"),
+    variable = all_variables
+  )
+
+  chains <- posterior::chain_ids(draws)
+  paths <- generate_file_names(basename = "fittedParams", ids = chains)
+  paths <- file.path(tempdir(), paths)
+  chain <- 1
+  for (path in paths) {
+    write(
+      paste0("# num_samples = ", n, "\n", paste0(unrepair_variable_names(all_variables), collapse = ",")),
+      file = path,
+      append = FALSE
+    )
+    utils::write.table(
+      posterior::subset_draws(draws, chain = chain),
+      sep = ",",
+      file = path,
+      col.names = FALSE,
+      row.names = FALSE,
+      append = TRUE
+    )
+    chain <- chain + 1
+  }
+  paths
+}
 
 #' Process fitted params for the generate quantities method
 #'
 #' @noRd
-#' @param fitted_params Paths to CSV files compatible with CmdStan or a CmdStanMCMC object.
+#' @param fitted_params Paths to CSV files produced by Cmdstan sampling,
+#'  a CmdStanMCMC or CmdStanVB object, a draws_array or draws_matrix.
 #' @return Paths to CSV files containing parameter values.
 #'
 process_fitted_params <- function(fitted_params) {
   if (is.character(fitted_params)) {
     paths <- absolute_path(fitted_params)
-  } else if (checkmate::test_r6(fitted_params, classes = ("CmdStanMCMC"))) {
-    if (all(file.exists(fitted_params$output_files()))) {
+  } else if (checkmate::test_r6(fitted_params, classes = "CmdStanMCMC") &&
+              all(file.exists(fitted_params$output_files()))) {
       paths <- absolute_path(fitted_params$output_files())
-    } else {
-      draws <- tryCatch(posterior::as_draws_array(fitted_params$draws()),
-        error=function(cond) {
-            stop("Unable to obtain draws from the fit (CmdStanMCMC) object.", call. = FALSE)
-        }
-      )
-      sampler_diagnostics <- tryCatch(posterior::as_draws_array(fitted_params$sampler_diagnostics()),
-        error=function(cond) {
-            stop("Unable to obtain sampler diagnostics from the fit (CmdStanMCMC) object.", call. = FALSE)
-        }
-      )
-      if (!is.null(draws)) {
-        variables <- posterior::variables(draws)
-        non_lp_variables <- variables[variables != "lp__"]
-        draws <- posterior::bind_draws(
-          posterior::subset_draws(draws, variable = "lp__"),
-          sampler_diagnostics,
-          posterior::subset_draws(draws, variable = non_lp_variables),
-          along = "variable"
-        )
-        variables <- posterior::variables(draws)
-        chains <- posterior::chain_ids(draws)
-        iterations <- posterior::niterations(draws)
-        paths <- generate_file_names(basename = "fittedParams", ids = chains)
-        paths <- file.path(tempdir(), paths)
-        chain <- 1
-        for (path in paths) {
-          chain_draws <- posterior::subset_draws(draws, chain = chain)
-          write(
-            paste0("# num_samples = ", iterations),
-            file = path
-          )
-          write(
-            paste0(unrepair_variable_names(variables), collapse = ","),
-            file = path,
-            append = TRUE
-          )
-          utils::write.table(
-            chain_draws,
-            file = path,
-            sep = ",",
-            col.names = FALSE,
-            row.names = FALSE,
-            append = TRUE
-          )
-          chain <- chain + 1
-        }
+  } else if(checkmate::test_r6(fitted_params, classes = c("CmdStanMCMC"))) {
+    draws <- tryCatch(fitted_params$draws(),
+      error=function(cond) {
+          stop("Unable to obtain draws from the fit object.", call. = FALSE)
       }
-    }
+    )
+    sampler_diagnostics <- tryCatch(fitted_params$sampler_diagnostics(),
+      error=function(cond) {
+          NULL
+      }
+    )
+    paths <- draws_to_csv(draws, sampler_diagnostics)
+  } else if(checkmate::test_r6(fitted_params, classes = c("CmdStanVB"))) {
+    draws <- tryCatch(fitted_params$draws(),
+      error=function(cond) {
+          stop("Unable to obtain draws from the fit object.", call. = FALSE)
+      }
+    )
+    paths <- draws_to_csv(posterior::as_draws_array(draws))
+  } else if (any(class(fitted_params) == "draws_array")){
+    paths <- draws_to_csv(fitted_params)
+  } else if (any(class(fitted_params) == "draws_matrix")){
+    paths <- draws_to_csv(posterior::as_draws_array(fitted_params))
   } else {
-    stop("'fitted_params' should be a vector of paths or a CmdStanMCMC object.", call. = FALSE)
+    stop(
+      "'fitted_params' must be a list of paths to CSV files, ",
+      "a CmdStanMCMC/CmdStanVB object, ",
+      "a posterior::draws_array or a posterior::draws_matrix.", call. = FALSE)
   }
   paths
 }

R/model.R

@@ -1430,9 +1430,10 @@ CmdStanModel$set("public", name = "variational", value = variational_method)
 #'
 #' @section Arguments:
 #'   * `fitted_params`: (multiple options) The parameter draws to use. One of the following:
-#'     - A [CmdStanMCMC] fitted model object.
-#'     - A character vector of paths to CmdStan CSV output files containing
-#'     parameter draws.
+#'     - A [CmdStanMCMC] or [CmdStanVB] fitted model object.
+#'     - A [posterior::draws_array] (for MCMC) or [posterior::draws_matrix] (for VB)
+#'       object returned by CmdStanR's [`$draws()`][fit-method-draws] method.
+#'     - A character vector of paths to CmdStan CSV output files.
 #'   * `data`, `seed`, `output_dir`, `parallel_chains`, `threads_per_chain`, `sig_figs`:
 #'   Same as for the [`$sample()`][model-method-sample] method.
 #'

tests/testthat/test-data.R

@@ -25,48 +25,32 @@ test_that("process_fitted_params() works with basic input types", {
 })
 
 test_that("process_fitted_params() errors with bad args", {
+  error_msg <- "'fitted_params' must be a list of paths to CSV files, a CmdStanMCMC/CmdStanVB object, a posterior::draws_array or a posterior::draws_matrix."
   expect_error(
     process_fitted_params(5),
-    "'fitted_params' should be a vector of paths or a CmdStanMCMC object."
+    error_msg
   )
   expect_error(
     process_fitted_params(NULL),
-    "'fitted_params' should be a vector of paths or a CmdStanMCMC object."
-  )
-  expect_error(
-    process_fitted_params(fit_vb),
-    "'fitted_params' should be a vector of paths or a CmdStanMCMC object."
+    error_msg
   )
   expect_error(
     process_fitted_params(fit_optimize),
-    "'fitted_params' should be a vector of paths or a CmdStanMCMC object."
-  )
-
-  fit_tmp <- testing_fit("bernoulli", method = "sample", seed = 123)
-  temp_file <- tempfile(fileext = ".rds")
-  saveRDS(fit_tmp, file = temp_file)
-  rm(fit_tmp)
-  gc()
-  fit_tmp_null <- readRDS(temp_file)
-  expect_error(
-    process_fitted_params(fit_tmp_null),
-    "Unable to obtain draws from the fit \\(CmdStanMCMC\\) object."
+    error_msg
   )
 
   fit_tmp <- testing_fit("bernoulli", method = "sample", seed = 123)
-  fit_tmp$draws()
   temp_file <- tempfile(fileext = ".rds")
   saveRDS(fit_tmp, file = temp_file)
   rm(fit_tmp)
   gc()
   fit_tmp_null <- readRDS(temp_file)
   expect_error(
     process_fitted_params(fit_tmp_null),
-    "Unable to obtain sampler diagnostics from the fit \\(CmdStanMCMC\\) object."
+    "Unable to obtain draws from the fit object."
   )
 })
 
-
 test_that("process_fitted_params() works if output_files in fit do not exist", {
   fit_ref <- testing_fit("bernoulli", method = "sample", seed = 123)
   fit_tmp <- testing_fit("bernoulli", method = "sample", seed = 123)
@@ -108,4 +92,118 @@ test_that("process_fitted_params() works if output_files in fit do not exist", {
   }
 })
 
+test_that("process_fitted_params() works with CmdStanMCMC", {
+  fit <- testing_fit("logistic", method = "sample", seed = 123)
+  fit_params_files <- process_fitted_params(fit)
+  expect_true(all(file.exists(fit_params_files)))
+  chain <- 1
+  for(file in fit_params_files) {
+    if (os_is_windows()) {
+      grep_path <- repair_path(Sys.which("grep.exe"))
+      fread_cmd <- paste0(grep_path, " -v '^#' --color=never ", file)
+    } else {
+      fread_cmd <- paste0("grep -v '^#' --color=never ", file)
+    }
+    suppressWarnings(
+      fit_params_tmp <- data.table::fread(
+        cmd = fread_cmd
+      )
+    )
+    fit_params_tmp <- posterior::as_draws_array(fit_params_tmp)
+    posterior::variables(fit_params_tmp) <- repair_variable_names(posterior::variables(fit_params_tmp))
+    expect_equal(
+      posterior::subset_draws(fit$draws(), variable = "lp__", chain = chain),
+      posterior::subset_draws(fit_params_tmp, variable = "lp__")
+    )
+    expect_equal(
+      posterior::subset_draws(fit$draws(), variable = c("alpha", "beta[1]", "beta[2]", "beta[3]"), chain = chain),
+      posterior::subset_draws(fit_params_tmp, variable = c("alpha", "beta[1]", "beta[2]", "beta[3]"),)
+    )
+    chain <- chain + 1
+  }
+})
 
+test_that("process_fitted_params() works with draws_array", {
+  fit <- testing_fit("logistic", method = "sample", seed = 123)
+  fit_params_files <- process_fitted_params(fit$draws())
+  expect_true(all(file.exists(fit_params_files)))
+  chain <- 1
+  for(file in fit_params_files) {
+    if (os_is_windows()) {
+      grep_path <- repair_path(Sys.which("grep.exe"))
+      fread_cmd <- paste0(grep_path, " -v '^#' --color=never ", file)
+    } else {
+      fread_cmd <- paste0("grep -v '^#' --color=never ", file)
+    }
+    suppressWarnings(
+      fit_params_tmp <- data.table::fread(
+        cmd = fread_cmd
+      )
+    )
+    fit_params_tmp <- posterior::as_draws_array(fit_params_tmp)
+    posterior::variables(fit_params_tmp) <- repair_variable_names(posterior::variables(fit_params_tmp))
+    expect_equal(
+      posterior::subset_draws(fit$draws(), variable = "lp__", chain = chain),
+      posterior::subset_draws(fit_params_tmp, variable = "lp__")
+    )
+    expect_equal(
+      posterior::subset_draws(fit$draws(), variable = c("alpha", "beta[1]", "beta[2]", "beta[3]"), chain = chain),
+      posterior::subset_draws(fit_params_tmp, variable = c("alpha", "beta[1]", "beta[2]", "beta[3]"),)
+    )
+    chain <- chain + 1
+  }
+})
+
+test_that("process_fitted_params() works with CmdStanVB", {
+  fit <- testing_fit("logistic", method = "variational", seed = 123)
+  file <- process_fitted_params(fit)
+  expect_true(file.exists(file))
+  if (os_is_windows()) {
+    grep_path <- repair_path(Sys.which("grep.exe"))
+    fread_cmd <- paste0(grep_path, " -v '^#' --color=never ", file)
+  } else {
+    fread_cmd <- paste0("grep -v '^#' --color=never ", file)
+  }
+  suppressWarnings(
+    fit_params_tmp <- data.table::fread(
+      cmd = fread_cmd
+    )
+  )
+  fit_params_tmp <- posterior::as_draws_array(fit_params_tmp)
+  posterior::variables(fit_params_tmp) <- repair_variable_names(posterior::variables(fit_params_tmp))
+  expect_equal(
+    posterior::subset_draws(posterior::as_draws_array(fit$draws()), variable = "lp__"),
+    posterior::subset_draws(fit_params_tmp, variable = "lp__")
+  )
+  expect_equal(
+    posterior::subset_draws(posterior::as_draws_array(fit$draws()), variable = c("alpha", "beta[1]", "beta[2]", "beta[3]")),
+    posterior::subset_draws(fit_params_tmp, variable = c("alpha", "beta[1]", "beta[2]", "beta[3]"))
+  )
+})
+
+test_that("process_fitted_params() works with draws_matrix", {
+  fit <- testing_fit("logistic", method = "variational", seed = 123)
+  file <- process_fitted_params(fit$draws())
+  expect_true(file.exists(file))
+  if (os_is_windows()) {
+    grep_path <- repair_path(Sys.which("grep.exe"))
+    fread_cmd <- paste0(grep_path, " -v '^#' --color=never ", file)
+  } else {
+    fread_cmd <- paste0("grep -v '^#' --color=never ", file)
+  }
+  suppressWarnings(
+    fit_params_tmp <- data.table::fread(
+      cmd = fread_cmd
+    )
+  )
+  fit_params_tmp <- posterior::as_draws_array(fit_params_tmp)
+  posterior::variables(fit_params_tmp) <- repair_variable_names(posterior::variables(fit_params_tmp))
+  expect_equal(
+    posterior::subset_draws(posterior::as_draws_array(fit$draws()), variable = "lp__"),
+    posterior::subset_draws(fit_params_tmp, variable = "lp__")
+  )
+  expect_equal(
+    posterior::subset_draws(posterior::as_draws_array(fit$draws()), variable = c("alpha", "beta[1]", "beta[2]", "beta[3]")),
+    posterior::subset_draws(fit_params_tmp, variable = c("alpha", "beta[1]", "beta[2]", "beta[3]"))
+  )
+})

tests/testthat/test-model-generate_quantities.R

@@ -70,3 +70,34 @@ test_that("generate_quantities work for different chains and parallel_chains", {
     fixed = TRUE
   )
 })
+
+test_that("generate_quantities works with draws_array", {
+  skip_on_cran()
+  fit_1_chain <- testing_fit("bernoulli", method = "sample", seed = 123, chains = 1)
+  expect_gq_output(
+    mod_gq$generate_quantities(data = data_list, fitted_params = fit_1_chain$draws())
+  )
+  expect_gq_output(
+    mod_gq$generate_quantities(data = data_list, fitted_params = fit$draws(), parallel_chains = 2)
+  )
+  expect_gq_output(
+    mod_gq$generate_quantities(data = data_list, fitted_params = fit$draws(), parallel_chains = 4)
+  )
+})
+
+fit <- testing_fit("bernoulli", method = "variational", seed = 123)
+mod_gq <- testing_model("bernoulli_ppc")
+data_list <- testing_data("bernoulli")
+fit_gq <- mod_gq$generate_quantities(data = data_list, fitted_params = fit)
+
+test_that("generate_quantities works with VB and draws_matrix", {
+  skip_on_cran()
+  fit <- testing_fit("bernoulli", method = "variational", seed = 123)
+  fit_gq <- mod_gq$generate_quantities(data = data_list, fitted_params = fit)
+  expect_gq_output(
+    mod_gq$generate_quantities(data = data_list, fitted_params = fit)
+  )
+  expect_gq_output(
+    mod_gq$generate_quantities(data = data_list, fitted_params = fit$draws())
+  )
+})