Survival Plot with ggplot

tidy_survfit.R

require(ggplot2)
require(dplyR)

# Convert Survfit object to tidy data frame. survival = TRUE means at x = 0 y = 1 (descending upon event)
# and FALSE means x = 0 y = 0 (ascending) 
tidy_km = function( survfit, strata_names = NULL, survival = FALSE) {
  #
  tidy_single_strata = function ( strata ) {
    print(strata$surv)
    data_frame(
      time     = strata$time,
      n.risk   = strata$n.risk,
      n.event  = strata$n.risk,
      n.censor = strata$n.censor, # number censored at given time
      surv     = switch(survival+1, strata$surv,  1 -strata$surv),
      upper    = switch(survival+1, strata$upper, 1-strata$lower),
      lower    = switch(survival+1, strata$lower, 1-strata$upper),
      std.err  = strata$n.censor
    )
  }

  # unstratified survfit
  if( is.null(survfit$strata) ) {
    tidy_single_strata( survfit )
  } else {
    if( !is.null(strata_names) &
        length(strata_names) != length(names( survfit$strata )) ) {
      stop( "strata names not correct length" )
    }
    if( is.null(strata_names) ) {
    }
    lapply( 1:length(strata_names), function(strat) {
      print(strata_names[strat])
      tidy_single_strata( survfit[strat] ) %>% mutate( group = strata_names[strat] )
    }) %>% bind_rows
  }
}

# Plots stratified survival curve, requires "group" variable to be present
plot_surv = function(tidy_km) {
  tidy_km %>%
    ggplot() +
    geom_ribbon(aes(ymin=lower, x=time, ymax=upper, fill=group), alpha=.075) +
    geom_step(aes(x=time, y=surv, colour=group), size=1) +
    geom_point(
      data = tidy_km %>% filter(n.censor > 0),
      aes(x=time, y=surv, colour=group),
      shape=3
    )
}

tidy_km(s, c("Low", "Moderate-Low", "Moderate-High", "High") %>% factor(., .), survival=TRUE) %>%
  plot_curve +
  xlab("Years") +
  ylab("% Recurrence") +
  scale_y_continuous(breaks=seq(0,1.1,by=0.1)) + 
  xlim(c(0,2))