Support all pp_check types for censored y

NEWS.md

@@ -24,6 +24,8 @@ thanks to Urs Kalbitzer. (#1280)
 `projpred`'s K-fold CV. (#1286)
 * Fix response values in `make_standata` for `bernoulli` families
 when only 1s are present thanks to Facundo Munoz. (#1298)
+* Fix `pp_check` for censored responses to work for all plot types
+thanks to Hayden Rabel. (#1327) 
 
 # brms 2.16.3
 

R/pp_check.R

@@ -150,21 +150,18 @@ pp_check.brmsfit <- function(object, type, ndraws = NULL, nsamples = NULL,
     object, newdata = newdata, resp = resp,
     re_formula = NA, check_response = TRUE, ...
   )
-  # censored responses are misleading when displayed in pp_check
-  bterms <- brmsterms(object$formula)
-  cens <- get_cens(bterms, data, resp = resp)
-  if (!is.null(cens)) {
-    warning2("Censored responses are not shown in 'pp_check'.")
-    take <- !cens
-    if (!any(take)) {
-      stop2("No non-censored responses found.")
+
+  # prepare plotting arguments
+  ppc_args <- nlist(y, yrep)
+  if (!is.null(group)) {
+    ppc_args$group <- data[[group]]
+  }
+  if (!is.null(x)) {
+    ppc_args$x <- data[[x]]
+    if (!is_like_factor(ppc_args$x)) {
+      ppc_args$x <- as.numeric(ppc_args$x)
     }
-    y <- y[take]
-    yrep <- yrep[, take, drop = FALSE]
   }
-  # most ... arguments are ment for the prediction function
-  for_pred <- names(dots) %in% names(formals(prepare_predictions.brmsfit))
-  ppc_args <- c(list(y, yrep), dots[!for_pred])
   if ("psis_object" %in% setdiff(names(formals(ppc_fun)), names(ppc_args))) {
     ppc_args$psis_object <- do_call(
       compute_loo, c(pred_args, criterion = "psis")
@@ -175,14 +172,38 @@ pp_check.brmsfit <- function(object, type, ndraws = NULL, nsamples = NULL,
       do_call(compute_loo, c(pred_args, criterion = "psis"))
     )
   }
-  if (!is.null(group)) {
-    ppc_args$group <- data[[group]]
-  }
-  if (!is.null(x)) {
-    ppc_args$x <- data[[x]]
-    if (!is_like_factor(ppc_args$x)) {
-      ppc_args$x <- as.numeric(ppc_args$x)
+
+  # censored responses are misleading when displayed in pp_check
+  bterms <- brmsterms(object$formula)
+  cens <- get_cens(bterms, data, resp = resp)
+  if (!is.null(cens) & type != 'km_overlay') {
+    warning2("Censored responses are not shown in 'pp_check'.")
+    take <- !cens
+    if (!any(take)) {
+      stop2("No non-censored responses found.")
+    }
+    ppc_args$y <- ppc_args$y[take]
+    ppc_args$yrep <- ppc_args$yrep[, take, drop = FALSE]
+    if (!is.null(ppc_args$group)) {
+      ppc_args$group <- ppc_args$group[take]
+    }
+    if (!is.null(ppc_args$x)) {
+      ppc_args$x <- ppc_args$x[take]
+    }
+    if (!is.null(ppc_args$psis_object)) {
+      # tidier to re-compute with subset
+      psis_args <- c(pred_args, criterion = "psis")
+      psis_args$newdata <- data[take, ]
+      ppc_args$psis_object <- do_call(compute_loo, psis_args)
+    }
+    if (!is.null(ppc_args$lw)) {
+      ppc_args$lw <- ppc_args$lw[,take]
     }
   }
+
+  # most ... arguments are meant for the prediction function
+  for_pred <- names(dots) %in% names(formals(prepare_predictions.brmsfit))
+  ppc_args <- c(ppc_args, dots[!for_pred])
+
   do_call(ppc_fun, ppc_args)
 }

tests/local/setup_local_tests.R

@@ -14,4 +14,6 @@ expect_range <- function(object, lower = -Inf, upper = Inf, ...) {
   testthat::expect_true(all(object >= lower & object <= upper), ...)
 }
 
+SW <- suppressWarnings
+
 context("local tests")

tests/local/tests.models_new.R

@@ -76,12 +76,27 @@ test_that("Survival model from brm doc works correctly", {
   me3 <- conditional_effects(fit3, method = "predict")
   expect_ggplot(plot(me3, ask = FALSE)[[2]])
   expect_range(LOO(fit3)$estimates[3, 1], 650, 740)
-
+
+  # posterior checks of the censored model
+  expect_ggplot(SW(pp_check(
+    fit3, type = 'dens_overlay_grouped', group = 'sex', ndraws = 10
+  )))
+  expect_ggplot(SW(pp_check(
+    fit3, type = 'intervals', x = 'patient', ndraws = NULL
+  )))
+  expect_ggplot(SW(pp_check(fit3, type = 'loo_intervals', ndraws = NULL)))
+  expect_ggplot(SW(pp_check(fit3, type = 'loo_pit_overlay', ndraws = 10)))
+
   # enables rstan specific functionality
   fit3 <- add_rstan_model(fit3)
   expect_range(LOO(fit3, moment_match = TRUE)$estimates[3, 1], 650, 740)
   bridge <- bridge_sampler(fit3)
   expect_true(is.numeric(bridge$logml))
+
+  testthat::skip_if_not_installed("ggfortify")
+  expect_ggplot(SW(pp_check(
+    fit3, type = 'km_overlay', status_y = kidney$censored, ndraws = 10
+  )))
 })
 
 test_that("Binomial model from brm doc works correctly", {