R/plot_survival.R

In dswatson/bioplotr: Pretty, simple, optionally interactive plots for bioinformatics analysis pipelines

#' Survival Curve(s)
#'
#' This function plots one or several survival curves for a given time-to-event
#' model.
#'
#' @param fit An object of class \code{\link[survival]{survfit}}.
#' @param fun An arbitrary function defining a transformation of the survival
#'   curve(s). Common transformations can be specified with a character
#'   argument, e.g. \code{"event"} for cumulative events, \code{"cumhaz"} for
#'   the cumulative hazard function, or \code{"pct"} for survival probability as
#'   a percentage.
#' @param CI Plot confidence intervals?
#' @param censor Include tick-marks to indicate censored subjects?
#' @param pal_curves String specifying the color palette to use when plotting
#'   multiple vectors. Options include \code{"ggplot"}, all qualitative color 
#'   schemes available in \code{\href{https://bit.ly/2ipuEjn}{RColorBrewer}}, 
#'   and the complete collection of \code{\href{http://bit.ly/2bxnuGB}{ggsci}} 
#'   palettes. Alternatively, a character vector of colors with length equal to 
#'   the number of strata in \code{fit}.
#' @param title Optional plot title.
#' @param leg.txt Optional legend title.
#' @param legend Legend position. Must be one of \code{"bottom"}, \code{"left"},
#'   \code{"top"}, \code{"right"}, \code{"bottomright"}, \code{"bottomleft"},
#'   \code{"topleft"}, or \code{"topright"}.
#' @param ... Additional arguments to be passed to \code{\link[survminer]{
#'   ggsurvplot}}.
#'
#' @details
#' Survival curves visualize the fit of a time-to-event model, e.g. a
#' Kaplan-Meier estimator or a Cox proportional hazards regression. They are an
#' essential tool in survival analysis, providing a simple and intuitive visual
#' summary of the relative risk associated with different conditions.
#'
#' \code{plot_survival} is a lightweight wrapper for the \code{
#' \link[survminer]{ggsurvplot}} function from the \code{survminer} package,
#' which offers an impressive array of customization options for plotting
#' survival curves. \code{plot_survival} alters the default output of \code{
#' ggsurvplot} to align it with that of other figures generated by \code{
#' bioplotr}, but you can override those settings and/or take advantage of extra
#' \code{ggsurvplot} functionalities by passing additional arguments to \code{
#' plot_survival}.
#'
#' @references
#' Andersen, P. & Gill, R. (1982).
#' \href{https://www.jstor.org/stable/2240714?seq=1#page_scan_tab_contents}{
#' Cox's regression model for counting processes, a large sample study}. \emph{
#' Annals of Statistics}, \strong{10}, 1100-1120.
#'
#' Cox, D.R. (1972).
#' \href{https://www.jstor.org/stable/2985181?seq=1#page_scan_tab_contents}{
#' Regression Models and Life-Tables}. \emph{J. R. Stat. Soc., Series B, 34}(2):
#' 187-220.
#'
#' Kaplan, E.L. & Meier, P. (1958).
#' \href{https://www.jstor.org/stable/2281868?seq=1#page_scan_tab_contents}{
#' Nonparametric estimation from incomplete observations}. \emph{J. Amer. Stat.
#' Assn., 53}(282): 457-481.
#'
#' @examples
#' library(survival)
#' fit <- survfit(Surv(time, status) ~ sex, data = lung)
#' plot_survival(fit)
#'
#' @export
#' @importFrom survminer ggsurvplot
#' @import dplyr
#' @import ggplot2
#'

plot_survival <- function(
  fit,
         fun = NULL,
          CI = FALSE,
      censor = TRUE,
  pal_curves = 'npg',
       title = ifelse(is.null(fit$strata), 'Survival Curve', 'Survival Curves'),
     leg.txt = NULL,
      legend = 'right', ...
) {
  # Preliminaries
  if(!fit %>% inherits('survfit')) {
    stop('fit must be an object of class survfit. Load the survival package ',
         'and see ?survfit for more details.')
  }
  if (fun %>% is.null) {
    ylab <- 'Survival Probability'
  } else if (fun == 'event') {
    ylab <- 'Cumulative Event'
  } else if (fun == 'cumhaz') {
    ylab <- 'Cumulative Hazard'
  } else if (fun == 'pct') {
    ylab <- 'Survival Probability (%)'
  }
  if (!fit$strata %>% is.null) {
    cols <- colorize(pal_curves, var_type = 'Categorical',
                     n = length(fit$strata))
  }
  if (leg.txt %>% is.null && !fit$strata %>% is.null) {
    leg.txt <- gsub('=.*', '', names(fit$strata))[1]
    leg.lbl <- gsub('.*=', '', names(fit$strata))
  } else if (!(leg.txt %>% is.null) && !(fit$strata %>% is.null)) {
    leg.lbl <- names(fit$strata)
  }
  locations <- c('bottom', 'left', 'top', 'right',
                 'bottomright', 'bottomleft', 'topleft', 'topright')
  legend <- match.arg(legend, locations)
  if (legend == 'bottomright') {
    legend <- c(0.99, 0.01)
  } else if (legend == 'bottomleft') {
    legend <- c(0.01, 0.01)
  } else if (legend == 'topleft') {
    legend <- c(0.01, 0.99)
  } else if (legend == 'topright') {
    legend <- c(0.99, 0.99)
  }

  # Build plot
  p <- ggsurvplot(fit, fun = fun, size = 0.5, conf.int = CI,
                  censor = censor, title = title, legend = legend,
                  font.tickslab = 9L, ggtheme = theme_bw(), ylab = ylab,
                  ...)$plot
  if (!fit$strata %>% is.null) {
    p <- p + scale_color_manual(name = leg.txt,
                              labels = leg.lbl,
                              values = cols)
  }
  p <- p + theme(plot.title = element_text(hjust = 0.5))

  # Output
  print(p)

}


# ggplotly? Would need to go over curves with a text = Curve aesthetic...
# Does this work with survival forests and other nonparametric models?
# Remove separate legend label for confidence intervals