Simulation-based calibration (SBC)

[wlandau]

8c27696 adds a new transform argument to tar_jags_rep_draws() to support simulation-based calibration (SBC). I do not have a section in the vignettes because the SBC R package is not yet on CRAN, but here is an example to demonstrate the use case. I use the model from the vignettes:

model {
  for (i in 1:n) {
    y[i] ~ dnorm(beta[1] + x[i] * beta[2], 1)
  }
  for (i in 1:2) {
    beta[i] ~ dnorm(0, 1)
  }
}

and the following _targets.R file:

# _targets.R file:
library(targets)
library(jagstargets)
options(crayon.enabled = FALSE)
tar_option_set(
  packages = c("dplyr", "posterior", "purrr", "tibble"),
  memory = "transient",
  garbage_collection = TRUE
)

generate_data <- function(n = 10L) {
  beta <- stats::rnorm(n = 2, mean = 0, sd = 1)
  x <- seq(from = -1, to = 1, length.out = n)
  y <- stats::rnorm(n, beta[1] + x * beta[2], 1)
  .join_data <- list(beta = beta) # Used by tar_jags_rep_summary() and get_ranks().
  list(n = n, x = x, y = y, .join_data = .join_data)
}

#' @title Get simulation-based calibration (SBC) ranks for a simulation rep.
#' @return A one-row tibble of SBC rank statistics
#'   with one column for each model parameter.
#' @param data JAGS data list.
#' @param draws Data frame with one row per MCMC sample and one column
#'   per model parameter
#' @examples
#' library(dplyr)
#' library(posterior)
#' library(purrr)
#' library(tibble)
#' data <- generate_data()
#' draws <- tibble(`beta[1]` = rnorm(100), `beta[2]` = rnorm(100))
#' get_ranks(data, draws)
get_ranks <- function(data, draws) {
  draws <- select(draws, starts_with(names(data$.join_data)))
  truth <- map_dbl(names(draws), ~eval(parse(text = .x), envir = data$.join_data))
  out <- SBC::calculate_ranks_draws_matrix(truth, as_draws_matrix(draws))
  as_tibble(as.list(out))
}

list(
  tar_jags_rep_draws(
    name = model,
    jags_files = "model.jags",
    data = generate_data(),
    parameters.to.save = "beta",
    batches = 5,
    reps = 4,
    transform = get_ranks
  )
)

After I run tar_make(), I get the following output for the rank statistics:

tar_read(model_model)
#> `beta[1]` `beta[2]` .dataset_id           .rep           .seed .file      .name
#> <dbl>     <dbl> <chr>                 <chr>          <int> <chr>      <chr>
#>   1      2443      2695 model_data_b0b9380a_1 3043ac7f -1490816250 model.jags model
#> 2      1976      1493 model_data_b0b9380a_2 31910c25  2036930449 model.jags model
#> 3      2369       190 model_data_b0b9380a_3 17939b8e   483483223 model.jags model
#> 4      1209      2456 model_data_b0b9380a_4 c1d48e85   664499179 model.jags model
#> 5      2213       716 model_data_ffcdb73c_1 f28ded83 -1333721664 model.jags model
#> 6      1801      2776 model_data_ffcdb73c_2 b8aee74d  -894289374 model.jags model
#> 7      1705      1626 model_data_ffcdb73c_3 cd36337c   -37052330 model.jags model
#> 8      2163      2683 model_data_ffcdb73c_4 29c59f6c  1333838785 model.jags model
#> 9      2719      2292 model_data_b968a03a_1 6ac2648b   573008143 model.jags model
#> 10      2631      1380 model_data_b968a03a_2 f54f9b0f  1825978703 model.jags model
#> 11      1220       307 model_data_b968a03a_3 d2cf87be -2126093777 model.jags model
#> 12       308      1628 model_data_b968a03a_4 fa53ac53  1213296710 model.jags model
#> 13       384       384 model_data_f8763cb2_1 92b6222f  1589550368 model.jags model
#> 14       799        43 model_data_f8763cb2_2 08e4b83a -1052483030 model.jags model
#> 15      1496      1532 model_data_f8763cb2_3 a6b2f352  -668758032 model.jags model
#> 16      2802      2552 model_data_f8763cb2_4 5690b746 -1630954187 model.jags model
#> 17      2934       532 model_data_0bfdabdc_1 818e2e94 -1394132624 model.jags model
#> 18       230      1318 model_data_0bfdabdc_2 c2a3ff73   141291648 model.jags model
#> 19       243      2330 model_data_0bfdabdc_3 fcf1704b   519642522 model.jags model
#> 20      2020      2819 model_data_0bfdabdc_4 a740850b  -273890098 model.jags model#>