R/plot_prepare.R

In csetraynor/rstanhaz: Bayesian Survival Models (rstanhaz)

get_colors <- function(n){
  colors <- c( "#89C5DA", "#DA5724", "#74D944", "#CE50CA", "#3F4921", "#C0717C", "#CBD588", "#5F7FC7", 
  "#673770", "#D3D93E", "#38333E", "#508578", "#D7C1B1", "#689030", "#AD6F3B", "#CD9BCD", 
  "#D14285", "#6DDE88", "#652926", "#7FDCC0", "#C84248", "#8569D5", "#5E738F", "#D1A33D", 
  "#8A7C64", "#599861")
  colors[1:n]
}

plot_bs <- function(brier_score, models){
  ### make beautiful plots
  
  names(brier_score) <- models
  data.plot <- lapply(seq_along(brier_score), function(d){
    data.frame(BS = brier_score[[d]]$bsc,
               time = brier_score[[d]]$time,
               model = names(brier_score[[d]]) 
               )
  })
  
  data.caption <- sapply(seq_along(brier_score), function(d){
    paste0("iBrier:\n ", models[[d]] ," = ", round( brier_score[[d]]$bsc.integrated, 3))
    })
  data.caption <- paste(data.caption, collapse = "\n")

  # The palette with black:
  cbbPalette <- get_colors(n = length(models))
  ggplot2::ggplot(data = do.call(rbind.data.frame, data.plot),
                        aes(x = time, y = BS)) +
    geom_line( aes(color = model), size = 0.55) +
    labs(title = "Prediction error curve\n", x = "Time [months]", y = "survival Brier score", color = "Model\n",
         caption = data.caption ) +
    scale_color_manual(labels = models, values = cbbPalette) +
    theme_bw() +
    theme(axis.text.x=element_text(size=14),
          axis.text.y=element_text(size=14),
          plot.title=element_text(size=15, face="bold", color="darkgreen") )
}

#' Prepare plot frame
#'
#' This function prepares plot frame
#'
#' @param post posterior draws in matrix format \cr
#' @param strata strata
#' @param obs original dataframe defaults to ic2surv
#' @return d a plot dataset
#' @export
#' @importFrom magrittr %>%
get_plot.frame <- function(post, strata, obs ){
  plot.frame <- tibble::tibble(
    postmean = apply(plot.matrix, 1, mean),
    lower = apply(plot.matrix, 1, quantile, probs = 0.055),
    upper = apply(plot.matrix, 1, quantile, probs = 0.945),
    time = obs %>% arrange(time) %>% select(time) %>% 
      unlist %>% as.double
  ) %>%
    cbind ( obs %>% 
              dplyr::arrange(time) %>% 
              dplyr::select(strata) )
  
  zeros.frame <- tibble::tibble(time=0, postmean=1, 
                            lower = 1, upper = 1) %>%
    cbind(plot.frame %>% select(strata) %>% unique)
  
  plot.frame <- rbind(plot.frame, zeros.frame)
  plot.frame
}

get_km.frame <- function(obs, strata, time = "time", status = "status"){
  form <- as.formula(paste0("Surv(", time, ",", status, ") ~ ", paste(strata, collapse = "+")))
  mle.surv <- survfit(form, data = obs  )
  obs.mortality <- data.frame(time = mle.surv$time,
                              surv = mle.surv$surv,
                              strata = summary(mle.surv, censored=T)$strata)
  
  zeros <- data.frame(time=0, surv=1, strata=unique(obs.mortality$strata))
  obs.mortality <- rbind(obs.mortality, zeros)
  
  strata <- str_strata(obs.mortality)
  obs.mortality$strata <- NULL
  obs.mortality <- cbind( obs.mortality, strata)
  obs.mortality 
}

#' Prepare plot with new dataframe
#'
#' This function prepares plot frame
#'
#' @param post posterior draws in matrix format \cr
#' @param strata strata
#' @param obs original dataframe defaults to ic2surv
#' @return d a plot dataset
#' @export
#' @importFrom magrittr %>%
get_plot_frame <- function(mod, x = NULL, long_x = NULL, treatment = NULL, time = "time", status = "status", unix = FALSE){
  if(is.null(long_x) ) {
    long_x <- gen_stan_dat(x, time = time, status = status)
  }
  survdata <- pred_surv(mod = mod, long_x = long_x, unix = unix)
  
  mean.surv <- apply(survdata, MARGIN = 1, mean)
  mean.upper <- apply(survdata, 1, quantile, 0.945)
  mean.lower <- apply(survdata, 1 , quantile, 0.055)
  
  bc.plot <- data.frame(
    survmean = c(1, mean.surv),
    survlower = c(1, mean.lower),
    survupper = c(1, mean.upper),
    time = c(0, long_x$time)
  ) 
  
  if(!is.null(treatment)) {
    bc.plot <-  bc.plot %>%
      dplyr::mutate(Treatment = treatment)
  }
  bc.plot
}