Add extrapolation_factor parameter to ppc_km_overlay and ppc_km_overlay_grouped

Author: Sakuski

R/ppc-censoring.R

@@ -88,7 +88,8 @@ ppc_km_overlay <- function(
   status_y,
   left_truncation_y = NULL,
   size = 0.25,
-  alpha = 0.7
+  alpha = 0.7,
+  extrapolation_factor = 1.2
 ) {
   check_ignored_arguments(..., ok_args = "add_group")
   add_group <- list(...)$add_group
@@ -106,6 +107,10 @@ ppc_km_overlay <- function(
     }
   }
 
+  if (extrapolation_factor < 1) {
+    stop("`extrapolation_factor` must be greater than or equal to 1.")
+  }
+
   data <- ppc_data(y, yrep, group = status_y)
 
   # Modify the status indicator:
@@ -149,6 +154,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)))