R/ggfitStrata.R

In christophergandrud/simtvc: Tools for simulating and graphing time-varying hazards.

#' Graph fitted stratified survival curves
#'
#' \code{ggfitStrata} graphs fitted survival curves created with \code{\link{survfit}} using \link{ggplot2}.
#'
#' @param obj a \code{survfit} object.
#' @param byStrata logical, whether or not you want to include all of the stratified survival curves on one plot or separate them into a grid arranged plot.
#' @param xlab a label for the plot's x-axis
#' @param ylab a label of the plot's y-axis
#' @param title plot title.
#' @param lcolour line color. Currently only works if \code{byStrata = TRUE}. The default it \code{lcolour = "#2C7FB8"} (a bluish hexadecimal colour)
#' @param rcolour confidence bounds ribbon color. Either a single color or a vector of colours. The default it \code{lcolour = "#2C7FB8"} (a bluish hexadecimal colour)
#'
#' @description This function largely improves \code{\link{plot.survfit}}. It plots the curves using \link{ggplot2} rather than base R graphics. One major advantage is the ability to split the survival curves into multiple plots and arrange them in a grid. This makes it easier to examine many strata at once. Otherwise they can be very bunched up.
#' 
#' Note: the strata legend labels need to be changed manually (see \code{\link{revalue}}) in the \code{survfit} object with the \code{strata} component.
#' @examples
#' # Load survival
#' library(survival)
#' # Subset data
#' bladder1 <- bladder[bladder$enum < 5, ] 
#' # Estimate coxph model
#' M1 <- coxph(Surv(stop, event) ~ (rx + size + number) * strata(enum) + cluster(id), bladder1)
#' # Survfit
#' M1Fit <- survfit(M1)
#' # Plot strata in a grid
#' ggfitStrata(M1Fit, byStrata = TRUE)
#'
#' @seealso \code{\link{survfit}}, \code{\link{ggplot2}} and \code{\link{strata}} 
#' @import ggplot2 gridExtra
#' @export 

ggfitStrata <- function(obj, byStrata = FALSE, xlab = "", ylab = "", title = "", lcolour = "#2C7FB8", rcolour = "#2C7FB8")
{
  require(ggplot2)
  sFit <- obj
  time <- sFit$time
  lower <- sFit$lower
  upper <- sFit$upper
  S <- sFit$surv
  strata <- sFit$strata
  strata <- factor(rep(names(strata), strata), levels = names(strata))
  TempData <- data.frame(Time = time, Lower = lower, 
                         Upper = upper, Survival = S, 
                         Strata = strata)   
  
  if (byStrata == FALSE){
    ggplot(data = TempData, aes(x = Time, 
                         y = Survival,
                         color = Strata,
                         fill = Strata)) +
      geom_line() +
      geom_ribbon(aes(ymin = Lower, ymax = Upper), alpha = I(0.1)) +
      xlab(xlab) + ylab(ylab) + ggtitle(title) +
      theme_bw()
      
  } else if (byStrata == TRUE){
    TempData$StrataC <- gsub("=", "", TempData$Strata)
    TempData$StrataC <- gsub(" ", "", TempData$StrataC)
    eachStrata <- unique(TempData$StrataC)
    p <- list()
    for (i in eachStrata){
      SubData <- subset(TempData, StrataC == i)
      p[[i]] <- ggplot(data = SubData, aes(x = Time, 
                                           y = Survival)) +
                                   geom_line(colour = lcolour) +
                                   geom_ribbon(aes(ymin = Lower, 
                                                ymax = Upper), 
                                                alpha = I(0.1),
                                                colour = NA,
                                                fill = rcolour) +
                                   xlab("") + ylab("") + ggtitle(paste(i, "\n")) +
                                   theme_bw()
    }
  Grid <- do.call(grid.arrange, c(p, main = title, sub = xlab, left = ylab))
  }
}