Add a new transform argument to support SBC

DESCRIPTION

@@ -58,17 +58,22 @@ Imports:
   withr (>= 2.1.2)
 Suggests:
   dplyr (>= 1.0.2),
+  ggplot2 (>= 3.0.0),
   knitr (>= 1.30),
+  purrr (>= 0.3.0),
   R.utils (>= 2.10.1),
   rmarkdown (>= 2.3),
+  SBC (>= 0.2.0),
   testthat (>= 3.0.0),
+  tidyr (>= 1.0.0),
   visNetwork (>= 2.0.9)
 Remotes:
+  hyunjimoon/SBC,
   stan-dev/cmdstanr,
 SystemRequirements: CmdStan >= 2.25.0
 Encoding: UTF-8
 Language: en-US
 Roxygen: list(markdown = TRUE)
-RoxygenNote: 7.2.1
+RoxygenNote: 7.2.3
 VignetteBuilder: knitr
 Config/testthat/edition: 3

R/tar_stan_gq_rep.R

@@ -45,6 +45,7 @@ tar_stan_gq_rep <- function(
   variables = NULL,
   summaries = NULL,
   summary_args = NULL,
+  transform = NULL,
   tidy_eval = targets::tar_option_get("tidy_eval"),
   packages = targets::tar_option_get("packages"),
   library = targets::tar_option_get("library"),
@@ -66,6 +67,7 @@ tar_stan_gq_rep <- function(
   targets::tar_assert_chr(stan_files)
   targets::tar_assert_unique(stan_files)
   lapply(stan_files, assert_stan_file)
+  assert_transform(transform)
   name_stan <- produce_stan_names(stan_files)
   name_file <- paste0(name, "_file")
   name_lines <- paste0(name, "_lines")
@@ -120,7 +122,8 @@ tar_stan_gq_rep <- function(
     data_copy = data_copy,
     variables = variables,
     summaries = summaries,
-    summary_args = summary_args
+    summary_args = summary_args,
+    transform = transform
   )
   command <- as.expression(as.call(args))
   pattern_data <- substitute(map(x), env = list(x = sym_batch))
@@ -275,7 +278,8 @@ tar_stan_gq_rep_run <- function(
   data_copy,
   variables,
   summaries,
-  summary_args
+  summary_args,
+  transform
 ) {
   if (!is.null(stdout)) {
     withr::local_output_sink(new = stdout, append = TRUE)
@@ -318,7 +322,8 @@ tar_stan_gq_rep_run <- function(
       data_copy = data_copy,
       variables = variables,
       summaries = summaries,
-      summary_args = summary_args
+      summary_args = summary_args,
+      transform = transform
     )
   )
   out$.file <- stan_path
@@ -340,7 +345,8 @@ tar_stan_gq_rep_run_rep <- function(
   data_copy,
   variables,
   summaries,
-  summary_args
+  summary_args,
+  transform
 ) {
   stan_seed <- data$.seed + 1L
   stan_seed <- if_any(is.null(seed), stan_seed, stan_seed + seed)
@@ -367,6 +373,7 @@ tar_stan_gq_rep_run_rep <- function(
     inc_warmup = NULL,
     data = data,
     data_copy = data_copy,
-    seed = stan_seed
+    seed = stan_seed,
+    transform = transform
   )
 }

R/tar_stan_gq_rep_draws.R

@@ -94,6 +94,7 @@ tar_stan_gq_rep_draws <- function(
   threads_per_chain = NULL,
   variables = NULL,
   data_copy = character(0),
+  transform = NULL,
   tidy_eval = targets::tar_option_get("tidy_eval"),
   packages = targets::tar_option_get("packages"),
   library = targets::tar_option_get("library"),
@@ -138,6 +139,7 @@ tar_stan_gq_rep_draws <- function(
     variables = variables,
     summaries = NULL,
     summary_args = NULL,
+    transform = substitute(transform),
     tidy_eval = tidy_eval,
     packages = packages,
     library = library,

R/tar_stan_mcmc_rep.R

@@ -93,6 +93,7 @@ tar_stan_mcmc_rep <- function(
   inc_warmup = FALSE,
   summaries = NULL,
   summary_args = NULL,
+  transform = NULL,
   tidy_eval = targets::tar_option_get("tidy_eval"),
   packages = targets::tar_option_get("packages"),
   library = targets::tar_option_get("library"),
@@ -115,6 +116,7 @@ tar_stan_mcmc_rep <- function(
   targets::tar_assert_unique(stan_files, "stan_files must be unique")
   targets::tar_assert_chr(data_copy, "data_copy must be a character vector")
   lapply(stan_files, assert_stan_file)
+  assert_transform(transform)
   name_stan <- produce_stan_names(stan_files)
   name_file <- paste0(name, "_file")
   name_lines <- paste0(name, "_lines")
@@ -193,7 +195,8 @@ tar_stan_mcmc_rep <- function(
     data_copy = data_copy,
     variables = variables,
     summaries = summaries,
-    summary_args = summary_args
+    summary_args = summary_args,
+    transform = transform
   )
   command <- as.expression(as.call(args))
   pattern_data <- substitute(map(x), env = list(x = sym_batch))
@@ -333,6 +336,12 @@ tar_stan_mcmc_rep <- function(
 #'   element of your Stan data list with names and dimensions corresponding
 #'   to those of the model. For details, read
 #'   <https://docs.ropensci.org/stantargets/articles/simulation.html>.
+#' @param transform Symbol or `NULL`, name of a function that accepts
+#'   arguments `data` and `draws` and returns a data frame. Here,
+#'   `data` is the JAGS data list supplied to the model, and `draws`
+#'   is a data frame with one column per model parameter and one row
+#'   per posterior sample. See the simulation-based calibration (SBC)
+#'   section of the simulation vignette for an example.
 tar_stan_mcmc_rep_run <- function(
   stan_file,
   stan_name,
@@ -382,7 +391,8 @@ tar_stan_mcmc_rep_run <- function(
   inc_warmup,
   variables,
   summaries,
-  summary_args
+  summary_args,
+  transform
 ) {
   if (!is.null(stdout)) {
     withr::local_output_sink(new = stdout, append = TRUE)
@@ -449,7 +459,8 @@ tar_stan_mcmc_rep_run <- function(
       inc_warmup = inc_warmup,
       variables = variables,
       summaries = summaries,
-      summary_args = summary_args
+      summary_args = summary_args,
+      transform = transform
     )
   )
   out$.file <- stan_path
@@ -495,7 +506,8 @@ tar_stan_mcmc_rep_run_rep <- function(
   inc_warmup,
   data_copy,
   summaries,
-  summary_args
+  summary_args,
+  transform
 ) {
   stan_seed <- data$.seed + 1L
   stan_seed <- if_any(is.null(seed), stan_seed, stan_seed + seed)
@@ -541,6 +553,7 @@ tar_stan_mcmc_rep_run_rep <- function(
     output_type = output_type,
     summaries = summaries,
     summary_args = summary_args,
+    transform = transform,
     variables = variables,
     inc_warmup = inc_warmup,
     data = data,

R/tar_stan_mcmc_rep_draws.R

@@ -6,7 +6,7 @@
 #' @inheritSection tar_stan_mcmc_rep Seeds
 #' @details Draws could take up a lot of storage. If storage becomes
 #'    excessive, please consider thinning the draws or using
-#'   `tar_stan_mcmc_rep_summaries()` instead.
+#'   `tar_stan_mcmc_rep_summary()` instead.
 #'
 #'   Most of the arguments are passed to the `$compile()`
 #'   and `$sample()` methods of the `CmdStanModel` class. If you
@@ -112,6 +112,7 @@ tar_stan_mcmc_rep_draws <- function(
   inc_warmup = FALSE,
   variables = NULL,
   data_copy = character(0),
+  transform = NULL,
   tidy_eval = targets::tar_option_get("tidy_eval"),
   packages = targets::tar_option_get("packages"),
   library = targets::tar_option_get("library"),
@@ -178,6 +179,7 @@ tar_stan_mcmc_rep_draws <- function(
     data_copy = data_copy,
     inc_warmup = inc_warmup,
     variables = variables,
+    transform = substitute(transform),
     tidy_eval = tidy_eval,
     packages = packages,
     library = library,

R/tar_stan_vb_rep.R

@@ -54,6 +54,7 @@ tar_stan_vb_rep <- function(
   variables = NULL,
   summaries = NULL,
   summary_args = NULL,
+  transform = NULL,
   tidy_eval = targets::tar_option_get("tidy_eval"),
   packages = targets::tar_option_get("packages"),
   library = targets::tar_option_get("library"),
@@ -75,6 +76,7 @@ tar_stan_vb_rep <- function(
   targets::tar_assert_chr(stan_files)
   targets::tar_assert_unique(stan_files)
   lapply(stan_files, assert_stan_file)
+  assert_transform(transform)
   name_stan <- produce_stan_names(stan_files)
   name_file <- paste0(name, "_file")
   name_lines <- paste0(name, "_lines")
@@ -139,7 +141,8 @@ tar_stan_vb_rep <- function(
     data_copy = data_copy,
     variables = variables,
     summaries = summaries,
-    summary_args = summary_args
+    summary_args = summary_args,
+    transform = transform
   )
   command <- as.expression(as.call(args))
   pattern_data <- substitute(map(x), env = list(x = sym_batch))
@@ -304,7 +307,8 @@ tar_stan_vb_rep_run <- function(
   data_copy,
   variables,
   summaries,
-  summary_args
+  summary_args,
+  transform
 ) {
   if (!is.null(stdout)) {
     withr::local_output_sink(new = stdout, append = TRUE)
@@ -357,7 +361,8 @@ tar_stan_vb_rep_run <- function(
       data_copy = data_copy,
       variables = variables,
       summaries = summaries,
-      summary_args = summary_args
+      summary_args = summary_args,
+      transform = transform
     )
   )
   out$.file <- stan_path
@@ -389,7 +394,8 @@ tar_stan_vb_rep_run_rep <- function(
   data_copy,
   variables,
   summaries,
-  summary_args
+  summary_args,
+  transform
 ) {
   stan_seed <- data$.seed + 1L
   stan_seed <- if_any(is.null(seed), stan_seed, stan_seed + seed)
@@ -426,6 +432,7 @@ tar_stan_vb_rep_run_rep <- function(
     inc_warmup = NULL,
     data = data,
     data_copy = data_copy,
-    seed = stan_seed
+    seed = stan_seed,
+    transform = transform
   )
 }

R/tar_stan_vb_rep_draws.R

@@ -99,6 +99,7 @@ tar_stan_vb_rep_draws <- function(
   sig_figs = NULL,
   data_copy = character(0),
   variables = NULL,
+  transform = NULL,
   tidy_eval = targets::tar_option_get("tidy_eval"),
   packages = targets::tar_option_get("packages"),
   library = targets::tar_option_get("library"),
@@ -153,6 +154,7 @@ tar_stan_vb_rep_draws <- function(
     variables = variables,
     summaries = NULL,
     summary_args = NULL,
+    transform = substitute(transform),
     tidy_eval = tidy_eval,
     packages = packages,
     library = library,

R/utils_assert.R

@@ -34,3 +34,26 @@ assert_inc_warmup_fit <- function(inc_warmup, inc_warmup_fit) {
     )
   }
 }
+
+assert_transform <- function(transform) {
+  if (is.null(transform)) {
+    return()
+  }
+  if (!is.symbol(transform)) {
+    targets::tar_throw_validate("transform must be a symbol or NULL.")
+  }
+  name <- as.character(transform)
+  msg <- "transform must be a function in the pipeline environment."
+  if (!(name %in% names(targets::tar_option_get("envir")))) {
+    targets::tar_throw_validate(msg)
+  }
+  fun <- targets::tar_option_get("envir")[[name]]
+  if (!is.function(fun)) {
+    targets::tar_throw_validate(msg)
+  }
+  args <- names(formals(fun))
+  if (!all(c("data", "draws") %in% args)) {
+    msg <- "transform must have arguments \"data\" and \"draws\""
+    targets::tar_throw_validate(msg)
+  }
+}

R/utils_output.R

@@ -14,6 +14,7 @@ tar_stan_output <- function(
   output_type,
   summaries,
   summary_args,
+  transform,
   variables,
   inc_warmup,
   data,
@@ -29,7 +30,13 @@ tar_stan_output <- function(
       summary_args = summary_args,
       variables = variables
     ),
-    draws = tar_stan_output_draws(fit, variables, inc_warmup),
+    draws = tar_stan_output_draws(
+      fit = fit,
+      data = data,
+      variables = variables,
+      inc_warmup = inc_warmup,
+      transform = transform
+    ),
     diagnostics = tar_stan_output_diagnostics(fit, inc_warmup)
   )
   out <- tibble::as_tibble(out)
@@ -57,13 +64,27 @@ tar_stan_output_summary <- function(
   eval(command)
 }
 
-tar_stan_output_draws <- function(fit, variables, inc_warmup) {
+tar_stan_output_draws <- function(
+  fit,
+  data,
+  variables,
+  inc_warmup,
+  transform
+) {
   out <- if_any(
     is.null(inc_warmup),
     fit$draws(variables = variables),
     fit$draws(variables = variables, inc_warmup = inc_warmup)
   )
-  tibble::as_tibble(posterior::as_draws_df(out))
+  out <- tibble::as_tibble(posterior::as_draws_df(out))
+  if (!is.null(transform)) {
+    out <- do.call(
+      what = transform,
+      args = list(data = data, draws = out),
+      envir = targets::tar_option_get("envir")
+    )
+  }
+  out
 }
 
 tar_stan_output_diagnostics <- function(fit, inc_warmup) {