Merge pull request #260 from fweber144/ppc_censoring_grouped

New function ppc_km_overlay_grouped()

Author: jgabry

NAMESPACE

@@ -120,6 +120,7 @@ export(ppc_intervals)
 export(ppc_intervals_data)
 export(ppc_intervals_grouped)
 export(ppc_km_overlay)
+export(ppc_km_overlay_grouped)
 export(ppc_loo_intervals)
 export(ppc_loo_pit)
 export(ppc_loo_pit_data)

NEWS.md

@@ -12,6 +12,9 @@
 * `mcmc_areas()` and `mcmc_areas_ridges()` gain an argument `border_size` for
 controlling the thickness of the ridgelines. (#224)
 
+* New plotting function `ppc_km_overlay_grouped()`, the grouped variant of
+  `ppc_km_overlay()`. (#260, @fweber144)
+
 # bayesplot 1.8.1
 
 * Fix R cmd check error on linux for CRAN

R/ppc-censoring.R

@@ -17,7 +17,7 @@
 #' @template args-y-yrep
 #' @param size,alpha Passed to the appropriate geom to control the appearance of
 #'   the `yrep` distributions.
-#' @param ... Currently unused.
+#' @param ... Currently only used internally.
 #'
 #' @template return-ggplot
 #'
@@ -30,6 +30,9 @@
 #'    `y`. Note that the replicated data from `yrep` is assumed to be
 #'    uncensored.
 #'   }
+#'   \item{`ppc_km_overlay_grouped()`}{
+#'    The same as `ppc_km_overlay()`, but with separate facets by `group`.
+#'   }
 #' }
 #'
 #' @templateVar bdaRef (Ch. 6)
@@ -50,81 +53,137 @@
 #' \donttest{
 #' ppc_km_overlay(y, yrep[1:25, ], status_y = status_y)
 #' }
+#' # With separate facets by group:
+#' group <- example_group_data()
+#' \donttest{
+#' ppc_km_overlay_grouped(y, yrep[1:25, ], group = group, status_y = status_y)
+#' }
 NULL
 
 #' @export
 #' @rdname PPC-censoring
 #' @param status_y The status indicator for the observations from `y`. This must
 #'   be a numeric vector of the same length as `y` with values in \{0, 1\} (0 =
 #'   right censored, 1 = event).
-ppc_km_overlay <-
-  function(y,
-           yrep,
-           ...,
-           status_y,
-           size = 0.25,
-           alpha = 0.7) {
-    check_ignored_arguments(...)
+ppc_km_overlay <- function(
+  y,
+  yrep,
+  ...,
+  status_y,
+  size = 0.25,
+  alpha = 0.7
+) {
+  check_ignored_arguments(..., ok_args = "add_group")
+  add_group <- list(...)$add_group
 
-    if(!requireNamespace("survival", quietly = TRUE)){
-      abort("Package 'survival' required.")
-    }
-    if(!requireNamespace("ggfortify", quietly = TRUE)){
-      abort("Package 'ggfortify' required.")
-    }
+  if(!requireNamespace("survival", quietly = TRUE)){
+    abort("Package 'survival' required.")
+  }
+  if(!requireNamespace("ggfortify", quietly = TRUE)){
+    abort("Package 'ggfortify' required.")
+  }
 
-    stopifnot(is.numeric(status_y))
-    stopifnot(all(status_y %in% c(0, 1)))
+  stopifnot(is.numeric(status_y))
+  stopifnot(all(status_y %in% c(0, 1)))
 
-    data <- ppc_data(y, yrep, group = status_y)
+  data <- ppc_data(y, yrep, group = status_y)
 
-    # Modify the status indicator:
-    #   * For the observed data ("y"), convert the status indicator back to
-    #     a numeric.
-    #   * For the replicated data ("yrep"), set the status indicator
-    #     to 1 ("event"). This way, the Kaplan-Meier estimator reduces
-    #     to "1 - ECDF" with ECDF denoting the ordinary empirical cumulative
-    #     distribution function.
-    data <- data %>%
-      dplyr::mutate(group = ifelse(.data$is_y,
-                                   as.numeric(as.character(.data$group)),
-                                   1))
+  # Modify the status indicator:
+  #   * For the observed data ("y"), convert the status indicator back to
+  #     a numeric.
+  #   * For the replicated data ("yrep"), set the status indicator
+  #     to 1 ("event"). This way, the Kaplan-Meier estimator reduces
+  #     to "1 - ECDF" with ECDF denoting the ordinary empirical cumulative
+  #     distribution function.
+  data <- data %>%
+    dplyr::mutate(group = ifelse(.data$is_y,
+                                 as.numeric(as.character(.data$group)),
+                                 1))
 
-    sf <- survival::survfit(
-      survival::Surv(value, group) ~ rep_label,
-      data = data
-    )
-    fsf <- fortify(sf)
+  sf_form <- survival::Surv(value, group) ~ rep_label
+  if(!is.null(add_group)){
+    data <- dplyr::inner_join(data,
+                              tibble::tibble(y_id = seq_along(y),
+                                             add_group = add_group),
+                              by = "y_id")
+    sf_form <- update(sf_form, . ~ . + add_group)
+  }
+  sf <- survival::survfit(
+    sf_form,
+    data = data
+  )
+  names(sf$strata) <- sub("add_group=", "add_group:", names(sf$strata)) # Needed to split the strata names in ggfortify:::fortify.survfit() properly.
+  fsf <- fortify(sf)
+  if(any(grepl("add_group", levels(fsf$strata)))){
+    strata_split <- strsplit(as.character(fsf$strata), split = ", add_group:")
+    fsf$strata <- as.factor(sapply(strata_split, "[[", 1))
+    fsf$group <- as.factor(sapply(strata_split, "[[", 2))
+  }
 
-    fsf$is_y_color <- as.factor(sub("\\[rep\\] \\(.*$", "rep", sub("^italic\\(y\\)", "y", fsf$strata)))
-    fsf$is_y_size <- ifelse(fsf$is_y_color == "yrep", size, 1)
-    fsf$is_y_alpha <- ifelse(fsf$is_y_color == "yrep", alpha, 1)
+  fsf$is_y_color <- as.factor(sub("\\[rep\\] \\(.*$", "rep", sub("^italic\\(y\\)", "y", fsf$strata)))
+  fsf$is_y_size <- ifelse(fsf$is_y_color == "yrep", size, 1)
+  fsf$is_y_alpha <- ifelse(fsf$is_y_color == "yrep", alpha, 1)
 
-    # 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)))
+  # 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)))
 
-    ggplot(data = fsf,
-           mapping = aes_(x = ~ time,
-                          y = ~ surv,
-                          color = ~ is_y_color,
-                          group = ~ strata,
-                          size = ~ is_y_size,
-                          alpha = ~ is_y_alpha)) +
-      geom_step() +
-      hline_at(
-        c(0, 0.5, 1),
-        size = c(0.2, 0.1, 0.2),
-        linetype = 2,
-        color = get_color("dh")
-      ) +
-      scale_size_identity() +
-      scale_alpha_identity() +
-      scale_color_ppc_dist() +
-      scale_y_continuous(breaks = c(0, 0.5, 1)) +
-      xlab(y_label()) +
-      yaxis_title(FALSE) +
-      xaxis_title(FALSE) +
-      yaxis_ticks(FALSE) +
-      bayesplot_theme_get()
-  }
+  ggplot(data = fsf,
+         mapping = aes_(x = ~ time,
+                        y = ~ surv,
+                        color = ~ is_y_color,
+                        group = ~ strata,
+                        size = ~ is_y_size,
+                        alpha = ~ is_y_alpha)) +
+    geom_step() +
+    hline_at(
+      0.5,
+      size = 0.1,
+      linetype = 2,
+      color = get_color("dh")
+    ) +
+    hline_at(
+      c(0, 1),
+      size = 0.2,
+      linetype = 2,
+      color = get_color("dh")
+    ) +
+    scale_size_identity() +
+    scale_alpha_identity() +
+    scale_color_ppc_dist() +
+    scale_y_continuous(breaks = c(0, 0.5, 1)) +
+    xlab(y_label()) +
+    yaxis_title(FALSE) +
+    xaxis_title(FALSE) +
+    yaxis_ticks(FALSE) +
+    bayesplot_theme_get()
+}
+
+#' @export
+#' @rdname PPC-censoring
+#' @template args-group
+ppc_km_overlay_grouped <- function(
+  y,
+  yrep,
+  group,
+  ...,
+  status_y,
+  size = 0.25,
+  alpha = 0.7
+) {
+  check_ignored_arguments(...)
+
+  p_overlay <- ppc_km_overlay(
+    y = y,
+    yrep = yrep,
+    add_group = group,
+    ...,
+    status_y = status_y,
+    size = size,
+    alpha = alpha
+  )
+
+  p_overlay +
+    facet_wrap("group") +
+    force_axes_in_facets()
+}

man-roxygen/args-group.R

@@ -1,3 +1,3 @@
-#' @param group A grouping variable (a vector or factor) the same length as
+#' @param group A grouping variable (a vector or factor) of the same length as
 #'   `y`. Each value in `group` is interpreted as the group level
 #'   pertaining to the corresponding value of `y`.