R/backend-mock.R

In hyunjimoon/SBC: Simulation Based Calibration for rstan/cmdstanr models

#' A mock backend.
#'
#' Mock backend is useful for testing the package.
#' It will ignore all data passed to it and always
#' provide `result` as the draws generated by the backend.
#'
#' @param result a `draws_matrix` that will be returned regardless of the data
#' @export
SBC_backend_mock <- function(result = posterior::draws_matrix(a = rnorm(100)),
                             output = NULL, message = NULL,
                             warning = NULL, error = NULL) {
  if(!posterior::is_draws_matrix(result)) {
    stop("Mock backend requires result to be draws_matrix")
  }

  structure(list(result = result,
                 output = output,
                 message = message,
                 warning = warning,
                 error = error), class = "SBC_backend_mock")
}

#' @export
SBC_fit.SBC_backend_mock <- function(backend, generated, cores) {
  if(!is.null(backend$output)) {
    cat(backend$output)
  }
  if(!is.null(backend$message)) {
    message(backend$message)
  }
  if(!is.null(backend$warning)) {
    warning(backend$warning)
  }
  if(!is.null(backend$error)) {
    stop(backend$error)
  }

  backend$result
}

#' @export
SBC_backend_mock_rng <- function(..., n_draws = 1000) {
  var_to_rng <- list(...)
  if(is.null(names(var_to_rng)) ||
     length(unique(names(var_to_rng))) != length(var_to_rng) ||
     any(names(var_to_rng) == "")
  ) {
    stop("All arguments must have a unique name")
  }
  var_to_rng <- purrr::map(var_to_rng, purrr::as_mapper)

  purrr::iwalk(var_to_rng, function(rng, name) {
    tryCatch({
      res <- rng(13)
    }, error = function(e) {
      message(e)
      stop("Test invocation for argument '", name, "' failed.\n",
           "All arguments must be convertible to a function that takes the number of draws as input and returns a 1D array of draws.")
    })
    if(!is.numeric(res) || length(res) != 13) {
      stop("Test invocation for argument '", name, "' returned unexpected result.\n",
           "All arguments must be convertible to a function that takes the number of draws as input and returns a 1D array of draws.")
    }
  })

  structure(list(var_to_rng = var_to_rng, n_draws = n_draws), class = "SBC_backend_mock_rng")
}

#' @export
SBC_fit.SBC_backend_mock_rng <- function(backend, generated, cores) {
  draws_list <- purrr::map(backend$var_to_rng, ~ .x(backend$n_draws))

  do.call(posterior::draws_matrix, draws_list)
}

#' @export
SBC_backend_iid_draws.SBC_backend_mock_rng <- function(backend) {
  TRUE
}