Add possibility to control extrapolation in ppc_km_overlay() and ppc_km_overlay_grouped()

NEWS.md

@@ -1,5 +1,6 @@
 # bayesplot (development version)
 
+* Add extrapolation_factor parameter to `ppc_km_overlay()` and `ppc_km_overlay_grouped()` by @Sakuski
 * Add possibility for left-truncation to `ppc_km_overlay()` and `ppc_km_overlay_grouped()` by @Sakuski
 * Added `ppc_loo_pit_ecdf()` by @TeemuSailynoja
 

R/ppc-censoring.R

@@ -24,11 +24,12 @@
 #' @section Plot Descriptions:
 #' \describe{
 #'   \item{`ppc_km_overlay()`}{
-#'    Empirical CCDF estimates of each dataset (row) in `yrep` are overlaid,
-#'    with the Kaplan-Meier estimate (Kaplan and Meier, 1958) for `y` itself on
-#'    top (and in a darker shade). This is a PPC suitable for right-censored
-#'    `y`. Note that the replicated data from `yrep` is assumed to be
-#'    uncensored.
+#'   Empirical CCDF estimates of each dataset (row) in `yrep` are overlaid, with
+#'   the Kaplan-Meier estimate (Kaplan and Meier, 1958) for `y` itself on top
+#'   (and in a darker shade). This is a PPC suitable for right-censored `y`.
+#'   Note that the replicated data from `yrep` is assumed to be uncensored. Left
+#'   truncation (delayed entry) times for `y` can be specified using
+#'   `left_truncation_y`.
 #'   }
 #'   \item{`ppc_km_overlay_grouped()`}{
 #'    The same as `ppc_km_overlay()`, but with separate facets by `group`.
@@ -40,24 +41,33 @@
 #' @template reference-km
 #'
 #' @examples
+#' \donttest{
 #' color_scheme_set("brightblue")
-#' y <- example_y_data()
+#'
 #' # For illustrative purposes, (right-)censor values y > 110:
+#' y <- example_y_data()
 #' status_y <- as.numeric(y <= 110)
 #' y <- pmin(y, 110)
+#'
 #' # In reality, the replicated data (yrep) would be obtained from a
 #' # model which takes the censoring of y properly into account. Here,
 #' # for illustrative purposes, we simply use example_yrep_draws():
 #' yrep <- example_yrep_draws()
 #' dim(yrep)
-#' \donttest{
+#'
+#' # Overlay 25 curves
 #' ppc_km_overlay(y, yrep[1:25, ], status_y = status_y)
-#' }
+#'
+#' # With extrapolation_factor = 1 (no extrapolation)
+#' ppc_km_overlay(y, yrep[1:25, ], status_y = status_y, extrapolation_factor = 1)
+#'
+#' # With extrapolation_factor = Inf (show all posterior predictive draws)
+#' ppc_km_overlay(y, yrep[1:25, ], status_y = status_y, extrapolation_factor = Inf)
+#'
 #' # With separate facets by group:
 #' group <- example_group_data()
-#' \donttest{
 #' ppc_km_overlay_grouped(y, yrep[1:25, ], group = group, status_y = status_y)
-#' }
+#'
 #' # With left-truncation (delayed entry) times:
 #' min_vals <- pmin(y, apply(yrep, 2, min))
 #' left_truncation_y <- rep(0, length(y))
@@ -66,7 +76,6 @@
 #'   runif(sum(condition), min = 0.6, max = 0.99) * y[condition],
 #'   min_vals[condition] - 0.001
 #' )
-#' \donttest{
 #' ppc_km_overlay(y, yrep[1:25, ], status_y = status_y,
 #'               left_truncation_y = left_truncation_y)
 #' }
@@ -78,15 +87,23 @@ NULL
 #'   be a numeric vector of the same length as `y` with values in \{0, 1\} (0 =
 #'   right censored, 1 = event).
 #' @param left_truncation_y Optional parameter that specifies left-truncation
-#'   (delayed entry) times for the observations from `y`. This must
-#'   be a numeric vector of the same length as `y`. If `NULL` (default),
-#'   no left-truncation is assumed.
+#'   (delayed entry) times for the observations from `y`. This must be a numeric
+#'   vector of the same length as `y`. If `NULL` (default), no left-truncation
+#'   is assumed.
+#' @param extrapolation_factor A numeric value (>=1) that controls how far the
+#'   plot is extended beyond the largest observed value in `y`. The default
+#'   value is 1.2, which corresponds to 20 % extrapolation. Note that all
+#'   posterior predictive draws may not be shown by default because of the
+#'   controlled extrapolation. To display all posterior predictive draws, set
+#'   `extrapolation_factor = Inf`.
+#'
 ppc_km_overlay <- function(
   y,
   yrep,
   ...,
   status_y,
   left_truncation_y = NULL,
+  extrapolation_factor = 1.2,
   size = 0.25,
   alpha = 0.7
 ) {
@@ -97,15 +114,25 @@ ppc_km_overlay <- function(
   suggested_package("ggfortify")
 
   if (!is.numeric(status_y) || length(status_y) != length(y) || !all(status_y %in% c(0, 1))) {
-    stop("`status_y` must be a numeric vector of 0s and 1s the same length as `y`.")
+    stop("`status_y` must be a numeric vector of 0s and 1s the same length as `y`.", call. = FALSE)
   }
 
   if (!is.null(left_truncation_y)) {
     if (!is.numeric(left_truncation_y) || length(left_truncation_y) != length(y)) {
-      stop("`left_truncation_y` must be a numeric vector of the same length as `y`.")
+      stop("`left_truncation_y` must be a numeric vector of the same length as `y`.", call. = FALSE)
     }
   }
 
+  if (extrapolation_factor < 1) {
+    stop("`extrapolation_factor` must be greater than or equal to 1.", call. = FALSE)
+  }
+  if (extrapolation_factor == 1.2) {
+    message(
+      "Note: `extrapolation_factor` now defaults to 1.2 (20%).\n",
+      "To display all posterior predictive draws, set `extrapolation_factor = Inf`."
+    )
+  }
+
   data <- ppc_data(y, yrep, group = status_y)
 
   # Modify the status indicator:
@@ -149,6 +176,10 @@ ppc_km_overlay <- function(
   fsf$is_y_size <- ifelse(fsf$is_y_color == "yrep", size, 1)
   fsf$is_y_alpha <- ifelse(fsf$is_y_color == "yrep", alpha, 1)
 
+  max_time_y <- max(y, na.rm = TRUE)
+  fsf <- fsf %>%
+    dplyr::filter(is_y_color != "yrep" | time <= max_time_y * extrapolation_factor)
+
   # Ensure that the observed data gets plotted last by reordering the
   # levels of the factor "strata"
   fsf$strata <- factor(fsf$strata, levels = rev(levels(fsf$strata)))
@@ -194,6 +225,7 @@ ppc_km_overlay_grouped <- function(
   ...,
   status_y,
   left_truncation_y = NULL,
+  extrapolation_factor = 1.2,
   size = 0.25,
   alpha = 0.7
 ) {
@@ -207,7 +239,8 @@ ppc_km_overlay_grouped <- function(
     status_y = status_y,
     left_truncation_y = left_truncation_y,
     size = size,
-    alpha = alpha
+    alpha = alpha,
+    extrapolation_factor = extrapolation_factor
   )
 
   p_overlay +