SBC implementations

[mike-lawrence]

@wlandau I see you've already proposed modifications to the posterior package to facilitate SBC, but I thought I'd point out this alternative (though admittedly less "official") implementation as possibly of interest.

[wlandau]

Thanks, it is encouraging to see this kind of work.

From the perspective of stantargets, I think the most expedient backup is just to manually compare the actual draws to .join_data to get the rank statistics. The actual calculation is trivially easy, but I hesitate because it would require new target factories for MCMC and VB, which sets a dangerous precedent for when many more alternative validation methods emerge.

[mike-lawrence]

I think the most expedient backup is just to manually compare the actual draws to .join_data to get the rank statistics.

At present, .join_data isn't appended to the draws tibble resulting from tar_stan_mcmc_rep_draws(), right?

[wlandau]

That's right. We would either need to append .join_data as a row or transform to long form before appending it as a column, both of which are contrary to how users expect the draws to be formatted. And we cannot append as an attribute because the draws are saved in compressed fst format, which only works on the simplest kinds of data frames.

In general, I want to encourage users to use summaries instead of draws when it comes to the *_rep_*() functions because draws take up a ton of data when repeated over thousands of reps.

[mike-lawrence]

Note I worked up an SBC demo using targets here:

https://github.com/mike-lawrence/sbc_demo

Possibly of interest for stantargets is how I stored the true parameters and then compute the ranks, plus added diagnostics. For those see "r_helpers/generate_true_and_data.r" and "r_helpers/sample_and_summarise.r"

[wlandau]

Thanks for sharing! I think I would be open to something like https://github.com/mike-lawrence/sbc_demo/blob/3505d6b2a31f00f0c4df533348ae556928c53a9b/r_helpers/sample_and_summarise.r#L15-L30 if stan-dev/posterior#111 does not work out.

(Aside: seems like we would probably only need this for MCMC, not sure I would trust variational Bayes for SBC.)

[martinmodrak]

Oh, I just found this discussion via stan-dev/posterior#111 - was thinking on whether integrating the SBC package (as linked by @mike-lawrence, https://github.com/hyunjimoon/SBC/) with (stan)targets would be sensible. Would be happy to hear your thoughts...

Since this discussion was opened, the package was almost completely rewritten and would probably now be much easier to integrate with targets. Some background on API is at https://hyunjimoon.github.io/SBC/articles/SBC.html and on how the "backends" work at https://hyunjimoon.github.io/SBC/articles/implementing_backends.html

To summarise:

1. User defines a "backend" (fitting algorithm + all parameters to the algorithm except data)
2. User defines a "generator" (generates simulated data)
3. The package calls the S3 generic SBC_fit to produce a fit
4. The fit is than transformed into all required outputs via: SBC_fit_to_draws_matrix.html(), SBC_statistics_from_single_fit() and SBC_fit_to_diagnostics()
5. Steps 3 & 4 are wrapped in a single internal function compute_SBC_single(), the main API entry compute_SBC() only handles caching, splits the work and then gathers the individual results in a single object.

The package aims to support also non-Stan models (currently only JAGS is built-in) and also supports both optimizing and VB fits in Stan (either via rstan or cmdstanr)

Currently we use future for parallelization, but that's a bit fragile with Stan models beyond a single-machine setup. To me the biggest problem in this was always figuring out how to only compile the model when necessary while ensuring correct propagation of the model to all workers (the current implementation in SBC just assumes all computation is done on the same machine). I was trying to read the docs/code here in stantargets but I am a bit unclear on how stantargets does this - it appears that if the model needs separate compilation at workers, the user needs to specify this manually via compile = "copy"? But then even if I have multiple workers per physical machine (but multiple physical machines), all workers will do their own compilation?

The reason I care about compilation is that I want to use the SBC package to throughout model development and early in the workflow, when one runs just a couple simulations to see if stuff is at least roughly OK, compilation times can be a substantial drag, so not repeating compilation when not absolutely needed seems important.

We also have a different approach to draws. In the SBC package, whole fits (and thus all draws) + all of stdout/messages/warnings/errors are kept by default (but there is an option to not transfer the fit from workers). The reason is once again workflow - I want to have the ability to investigate the fits and better understand where any problems arise (see e.g. https://hyunjimoon.github.io/SBC/articles/small_model_workflow.html for a workflow as I use it currently, where inspecting the fits is relevant). A possible alternative is to store the fits on disk and keep the ability to retrieve them on-demand, but not actually transfer them...

[wlandau]

Great work on the SBC package. For inclusion in stantargets, one would define a new output type to take a posterior::as_draws_df() object and produce a single data frame of ranks and/or ecdf info. Then a new tar_stan_mcmc_rep_sbc() target factory would wrap tar_stan_mcmc_rep() for that output type. stantargets takes care of the parallel computing/workflow stuff, and it is up to the user to write an R function to simulate a dataset from the model. So if SBC makes it easy to go from posterior::as_draws_df() to ranks/ecdf in a quick one-liner, then I will add it to stantargets, jagstargets, and other Bayesian Targetopia packages that may come up.

compile = "original" compiles the model once and then assumes all the workers that run the downstream sampling have access to that model. But on some platforms, the workers may not have access to the original file system. In that case, one can set compile = "copy" and each batch of model runs will compile its own copy of the Stan model. targets does not know what architecture you have, it only knows how you have batched the model runs into dynamic branches, so it makes each branch/batch do its own compilation.

[dmi3kno]

I guess you know it (from discussion in stan-dev/posterior#111), but the function calculate_ranks_draws_matrix() (docs: https://hyunjimoon.github.io/SBC/reference/calculate_ranks_draws_matrix.html, code: https://github.com/hyunjimoon/SBC/blob/master/R/results.R#L1022) does the calculation of the ranks for given variable values. It is not in df format, but that should not matter much. This one does strict "<=" comparison (which would correspond to Type 1 in Hyndman & Fan(1996)? ) with stochastic rank conflict resolution and therefore may not correspond directly to the default algorithm in posterior::quantile2() (which is Type 7).

References:
Hyndman RJ, Fan Y. 1996. Sample Quantiles in Statistical Packages. The American Statistician. 50(4):361–365. https://doi.org/10.2307/2684934

[wlandau]

With a custom target, I think you could call calculate_ranks_draws_matrix() on the output of a tar_stan_mcmc_rep_draws(). This approach requires saving all the posterior draws, which could get cumbersome if the model is high-dimensional or requires many iterations for convergence, but I think it is doable. I tried it with the small Stan model from the vignettes:

data {
  int <lower = 1> n;
  vector[n] x;
  vector[n] y;
}
parameters {
  vector[2] beta;
}
model {
  y ~ normal(beta[1] + x * beta[2], 1);
  beta ~ normal(0, 1);
}

Then in the _targets.R file, I added a new sbc target and new functions get_ranks_batch() and get_ranks_rep():

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

#' @title Generate a Stan data list for the model.
#' @return A Stan data list.
#' @param n Number of observations.
#' @examples
#' generate_data()
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_stan_mcmc_rep_summary()
  list(n = n, x = x, y = y, .join_data = .join_data)
}

#' @title Get the rank statistics of a batch of reps.
#' @return A tibble of rank statistics with one row per simulation rep
#'   and one column per model parameter.
#' @param data A list of Stan datasets produced by `generate_data()`
#' @param draws A tibble of draws with the `.rep` column to distinguish among the reps.
#' @param names A character vector of parameter names to get rank statistics.
#' @examples
#' library(dplyr)
#' library(posterior)
#' library(purrr)
#' library(tibble)
#' library(tidyr)
#' reps <- 4 # reps in the batch
#' data <- replicate(reps, generate_data(), simplify = FALSE)
#' draws <- tibble(`beta[1]` = rnorm(1e3), `beta[2]` = rnorm(1e3)) %>%
#'   expand_grid(.rep = map_chr(seq_len(reps), digest, algo = "xxhash32"))
#' names <- c("beta[1]", "beta[2]")
#' get_ranks_batch(data, draws, names)
get_ranks_batch <- function(data, draws, names) {
  draws_list <- split(draws, draws$.rep)[unique(draws$.rep)]
  map_dfr(seq_along(data), ~get_ranks_rep(data[[.x]], draws_list[[.x]], names))
}

#' @title Get the rank statistics of a rep.
#' @return A one-row tibble of rank statistics with
#'   one column per model parameter.
#' @param data A Stan data list from `generate_data()`
#' @param draws A tibble of posterior draws.
#' @param names A character vector of parameter names to get rank statistics.
#' @examples
#' library(dplyr)
#' library(posterior)
#' library(purrr)
#' library(tibble)
#' data <- generate_data()
#' draws <- tibble(`beta[1]` = rnorm(1e3), `beta[2]` = rnorm(1e3))
#' names <- c("beta[1]", "beta[2]")
#' get_ranks_rep(data, draws, names)
get_ranks_rep <- function(data, draws, names) {
  truth <- map_dbl(names, ~eval(parse(text = .x), envir = data$.join_data))
  draws_matrix <- as_draws_matrix(select(draws, all_of(names)))
  out <- SBC::calculate_ranks_draws_matrix(truth, draws_matrix)
  as_tibble(as.list(out))
}

list(
  tar_stan_mcmc_rep_draws(
    name = model,
    stan_files = "model.stan",
    data = generate_data(),
    batches = 5,
    reps = 4
  ),
  tar_target(
    name = sbc,
    command = get_ranks_batch(model_data, model_model, c("beta[1]", "beta[2]")),
    pattern = map(model_data, model_model),
    format = "fst_tbl"
  )
)

I ran the pipeline, and the result looks like this:

tar_read(sbc)
#> # A tibble: 20 × 2
#>    `beta[1]` `beta[2]`
#>        <dbl>     <dbl>
#>  1      2600       658
#>  2      2616      2049
#>  3      3098       298
#>  4      1646      3284
#>  5      1122      3083
 #> 6       802      2402
 #> 7      3083       432
 #> 8      2835      3582
#>  9      3647      3022
#> 10      3543      1781
#> 11      1869      1261
#> 12       433      2192
#> 13       548       529
#> 14      2147      3889
#> 15      3264      1553
#> 16      3464      2038
#> 17       177      3240
#> 18       303      1770
#> 19       339      3098
#> 20      3229      3491

[wlandau]

In f37eb66, I just added a transform argument to tar_stan_mcmc_rep_summary() which allows you to supply your own function to do SBC. This is a better approach because it's simpler (you don't need to disentangle targets batches) and you no longer need to save all the posterior samples of every simulation rep. I added a section of the simulation vignette to explain. Should publish to https://docs.ropensci.org/stantargets/articles/simulation.html#simulation-based-calibration when the GitHub Actions on the rOpenSci r-universe complete.

stantargets/vignettes/simulation.Rmd

Lines 392 to 560 in 942f0cf

	## Simulation-based calibration
	This section explores a more rigorous validation study which adopts the proper simulation-based calibration (SBC) method from [@talts2020]. To use this method, we need a function that generates rank statistics from a simulated dataset and a data frame of posterior draws. If the model is implemented correctly, these rank statistics will be uniformly distributed for each model parameter. Our function will use the [`calculate_ranks_draws_matrix()`](https://hyunjimoon.github.io/SBC/reference/calculate_ranks_draws_matrix.html) function from the [`SBC`](https://hyunjimoon.github.io/SBC/) R package [@sbc].
	```{r, eval = TRUE}
	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))
	}
	```
	To demonstrate this function, we simulate a dataset,
	```{r, eval = TRUE}
	data <- simulate_data()
	str(data)
	```
	we make up a hypothetical set of posterior draws,
	```{r, eval = TRUE}
	draws <- tibble(`beta[1]` = rnorm(100), `beta[2]` = rnorm(100))
	draws
	```
	and we call `get_ranks()` to get the SBC rank statistics for each model parameter.
	```{r, eval = TRUE}
	library(dplyr)
	library(posterior)
	library(purrr)
	get_ranks(data = data, draws = draws)
	```
	To put this into practice in a pipeline, we supply the symbol `get_ranks` to the `transform` argument of `tar_stan_mcmc_rep_draws()`. That way, instead of a full set of draws, each replication will return only the output of `get_ranks()` on those draws (plus a few helper columns). If supplied, the `transform` argument of `tar_stan_mcmc_rep_draws()` must be the name of a function in the pipeline. This function must accept arguments `data` and `draws`, and it must return a data frame.
	```{r, echo = FALSE}
	library(targets)
	tar_script({
	library(stantargets)
	options(crayon.enabled = FALSE)
	tar_option_set(
	packages = c("dplyr", "posterior", "purrr", "tibble"),
	memory = "transient",
	garbage_collection = TRUE
	)
	simulate_data <- function(n = 10L) {
	beta <- rnorm(n = 2, mean = 0, sd = 1)
	x <- seq(from = -1, to = 1, length.out = n)
	y <- rnorm(n, beta[1] + x * beta[2], 1)
	list(
	n = n,
	x = x,
	y = y,
	.join_data = list(beta = beta)
	)
	}
	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_stan_mcmc_rep_draws(
	model,
	"model.stan",
	simulate_data(),
	batches = 5,
	reps = 2,
	variables = "beta",
	stdout = R.utils::nullfile(),
	stderr = R.utils::nullfile(),
	transform = get_ranks
	)
	)
	})
	```
	```{r, eval = FALSE}
	# _targets.R
	library(targets)
	library(stantargets)
	tar_option_set(packages = c("dplyr", "posterior", "purrr", "tibble"))
	simulate_data <- function(n = 10L) {
	beta <- rnorm(n = 2, mean = 0, sd = 1)
	x <- seq(from = -1, to = 1, length.out = n)
	y <- rnorm(n, beta[1] + x * beta[2], 1)
	list(
	n = n,
	x = x,
	y = y,
	.join_data = list(beta = beta)
	)
	}
	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_stan_mcmc_rep_draws(
	model,
	c("model.stan"),
	simulate_data(),
	batches = 5,
	reps = 2,
	variables = "beta",
	stdout = R.utils::nullfile(),
	stderr = R.utils::nullfile(),
	transform = get_ranks # Supply the transform to get SBC ranks.
	)
	)
	```
	Our new function `get_ranks()` is a dependency of one of our downstream targets, so any changes to `get_ranks()` will force the results to refresh in the next run of the pipeline.
	```{r, eval = TRUE}
	tar_visnetwork()
	```
	Let's run the pipeline to compute a set of rank statistics for each simulated dataset.
	```{r, eval = TRUE, output = FALSE, warning = FALSE}
	tar_make()
	```
	We have a data frame of rank statistics with one row per simulation rep and one column per model parameter.
	```{r, eval = TRUE}
	tar_load(model_model)
	model_model
	```
	If the model is implemented correctly, then each the rank statistics each model parameter should be uniformly distributed. In practice, you may need thousands of simulation reps to make a judgment.
	```{r, eval = TRUE}
	library(ggplot2)
	library(tidyr)
	model_model %>%
	pivot_longer(
	starts_with("beta"),
	names_to = "parameter",
	values_to = "ranks"
	) %>%
	ggplot(.) +
	geom_histogram(aes(x = ranks), bins = 10) +
	facet_wrap(~parameter) +
	theme_gray(16)
	```