related to "proof of concept for posterior_pit" PR in brms

R/posterior_predict.R

@@ -82,7 +82,7 @@
 posterior_predict.brmsfit <- function(
   object, newdata = NULL, re_formula = NULL, re.form = NULL,
   transform = NULL, resp = NULL, negative_rt = FALSE,
-  ndraws = NULL, draw_ids = NULL, sort = FALSE, ntrys = 5, type = "r",
+  ndraws = NULL, draw_ids = NULL, sort = FALSE, ntrys = 5, output = "random",
   cores = NULL, ...
 ) {
   cl <- match.call()
@@ -93,11 +93,11 @@ posterior_predict.brmsfit <- function(
   object <- restructure(object)
   prep <- prepare_predictions(
     object, newdata = newdata, re_formula = re_formula, resp = resp,
-    ndraws = ndraws, draw_ids = draw_ids, check_response = FALSE, type = type, ...
+    ndraws = ndraws, draw_ids = draw_ids, check_response = FALSE, ...
   )
   posterior_predict(
     prep, transform = transform, sort = sort, ntrys = ntrys,
-    negative_rt = negative_rt, cores = cores, summary = FALSE, type = type
+    negative_rt = negative_rt, cores = cores, summary = FALSE, output = output
   )
 }
 
@@ -119,7 +119,7 @@ posterior_predict.mvbrmsprep <- function(object, ...) {
 posterior_predict.brmsprep <- function(object, transform = NULL, sort = FALSE,
                                        summary = FALSE, robust = FALSE,
                                        probs = c(0.025, 0.975),
-                                       cores = NULL, type = "r", ...) {
+                                       cores = NULL, ...) {
   summary <- as_one_logical(summary)
   cores <- validate_cores_post_processing(cores)
   if (is.customfamily(object$family)) {
@@ -136,7 +136,7 @@ posterior_predict.brmsprep <- function(object, transform = NULL, sort = FALSE,
   pp_fun <- paste0("posterior_predict_", object$family$fun)
   pp_fun <- get(pp_fun, asNamespace("brms"))
   N <- choose_N(object)
-  out <- plapply(seq_len(N), pp_fun, .cores = cores, prep = object, type = type, ...)
+  out <- plapply(seq_len(N), pp_fun, .cores = cores, prep = object, output = output, ...)
   if (grepl("_mv$", object$family$fun)) {
     out <- do_call(abind, c(out, along = 3))
     out <- aperm(out, perm = c(1, 3, 2))
@@ -303,50 +303,32 @@ validate_pp_method <- function(method) {
 
 # ------------------- family specific posterior_predict methods ---------------------
 # All posterior_predict_<family> functions have the same arguments structure
-# @param i index of the observatio for which to compute pp values
+# @param i index of the observation for which to compute pp values
 # @param prep A named list returned by prepare_predictions containing
 #   all required data and posterior draws
 # @param ... ignored arguments
 # @param A vector of length prep$ndraws containing draws
 #   from the posterior predictive distribution
-posterior_predict_gaussian <- function(i, prep, ntrys = 5, type = "r", ...) {
+posterior_predict_gaussian <- function(i, prep, ntrys = 5, output = "random", ...) {
   mu <- get_dpar(prep, "mu", i = i)
   sigma <- get_dpar(prep, "sigma", i = i)
   sigma <- add_sigma_se(sigma, prep, i = i)
-  switch(type,
-         r = rcontinuous(
-           n = prep$ndraws, dist = "norm",
-           mean = mu, sd = sigma,
-           lb = prep$data$lb[i], ub = prep$data$ub[i],
-           ntrys = ntrys
-         ),
-         p = pcontinuous(
-           n = prep$ndraws, dist = "norm",
-           q = prep$data$Y[i], mean = mu, sd = sigma,
-           lb = prep$data$lb[i], ub = prep$data$ub[i],
-           ntrys = ntrys
-         )
+
+  .predict_continuous_helper(
+    output = output, prep = prep, i = i, ntrys = ntrys,
+    dist = "norm", mean = mu, sd = sigma, ...
   )
 }
 
-posterior_predict_student <- function(i, prep, ntrys = 5, type = "r", ...) {
+posterior_predict_student <- function(i, prep, ntrys = 5, output = "random", ...) {
   nu <- get_dpar(prep, "nu", i = i)
   mu <- get_dpar(prep, "mu", i = i)
   sigma <- get_dpar(prep, "sigma", i = i)
   sigma <- add_sigma_se(sigma, prep, i = i)
-  switch(type,
-         r =   rcontinuous(
-           n = prep$ndraws, dist = "student_t",
-           df = nu, mu = mu, sigma = sigma,
-           lb = prep$data$lb[i], ub = prep$data$ub[i],
-           ntrys = ntrys
-         ),
-         p = pcontinuous(
-           n = prep$ndraws, dist = "student_t",
-           q = prep$data$Y[i], df = nu, mu = mu, sigma = sigma,
-           lb = prep$data$lb[i], ub = prep$data$ub[i],
-           ntrys = ntrys
-         )
+
+  .predict_continuous_helper(
+    output = output, prep = prep, i = i, ntrys = ntrys,
+    dist = "student_t", df = nu, mu = mu, sigma = sigma, ...
   )
 }
 
@@ -1018,12 +1000,12 @@ rcontinuous <- function(n, dist, ..., lb = NULL, ub = NULL, ntrys = 5) {
   out
 }
 
-pcontinuous <- function(n, dist, ..., lb = NULL, ub = NULL, ntrys = 5) {
+pcontinuous <- function(q, dist, ..., lb = NULL, ub = NULL, ntrys = 5) {
   args <- list(...)
   if (is.null(lb) && is.null(ub)) {
     # sample as usual
     pdist <- paste0("p", dist)
-    out <- do_call(pdist, c(list(n), args))
+    out <- do_call(pdist, c(list(q), args))
   } else {
     error("not implemented yet")
     # sample from truncated distribution
@@ -1129,3 +1111,34 @@ check_discrete_trunc_bounds <- function(x, lb = NULL, ub = NULL, thres = 0.01) {
   }
   round(x)
 }
+
+# predict random numbers or probability values from continuous distributions
+# @param output "probability" or "random"
+# @param prep A named list returned by prepare_predictions containing
+#   all required data and posterior draws
+# @param i index of the observation for which to compute pp values
+# @param dist name of the distribution
+# @param ntrys number of trys in rejection sampling for truncated models
+# @param q optional custom quantile value; if NULL, the default is prep$data$Y[i] 
+# @param ... additional arguments passed to the distribution functions
+# @return a vector of draws
+.predict_continuous_helper <- function(output, prep, i, dist, ntrys, q = NULL, ...) {
+  lb <- prep$data$lb[i]
+  ub <- prep$data$ub[i]
+
+  switch(output,
+    "probability" = {
+      if (is.null(q)) {
+        q <- prep$data$Y[i]
+      }
+      pcontinuous(
+        q = q, dist = dist, lb = lb, ub = ub, ntrys = ntrys, ...
+      )
+    },
+    "random" = {
+      rcontinuous(
+        n = prep$ndraws, dist = dist, lb = lb, ub = ub, ntrys = ntrys, ...
+      )
+    }
+  )
+}

tests/testthat/tests.posterior_predict.R

@@ -439,3 +439,26 @@ test_that("posterior_predict_custom runs without errors", {
   }
   expect_equal(length(brms:::posterior_predict_custom(sample(1:nobs, 1), prep)), ns)
 })
+
+test_that("posterior_predict_gaussian runs with various 'output' values without error", {
+  fit <- rename_pars(brms:::brmsfit_example3)
+  prep <- brms::prepare_predictions(fit)
+  model_fit <- fit$fit@sim
+  S <- model_fit$chains * (model_fit$iter - model_fit$warmup)
+  i <- 1
+
+  # probability
+  rpred <- brms:::posterior_predict_gaussian(i, prep = prep, output = "random")
+  expect_equal(length(rpred), S)
+
+  # compute PIT values (q = prep$data$Y[i])
+  PITs <- brms:::posterior_predict_gaussian(i, prep = prep, output = "probability")
+  expect_equal(length(PITs), S)
+  expect_true(all(PITs >= 0 & PITs <= 1))
+
+  # compute cdf based on custom 'q'
+  qpred <- brms:::posterior_predict_gaussian(i, q = 15, prep = prep, output = "probability")
+  expect_equal(length(qpred), S)
+  expect_true(all(qpred >= 0 & qpred <= 1))
+})
+