R/compute_kaplan.R

In ronkeizer/vpc: Create Visual Predictive Checks

#' Compute Kaplan-Meier statistics
#'
#' @param dat data.frame with events
#' @param strat vector of stratification variables
#' @param reverse_prob reverse the probability (i.e. return `1-probability`)?
#' @param rtte_conditional `TRUE` (default) or `FALSE`. Compute the probability for each event newly (`TRUE`), or calculate the absolute probability (`FALSE`, i.e. the "probability of a 1st, 2nd, 3rd event etc" rather than the "probability of an event happening").
#' @param ci confidence interval to calculate, numeric vector of length 2
compute_kaplan <- function(dat,
                           strat = "strat",
                           reverse_prob = FALSE,
                           rtte_conditional = TRUE,
                           ci = NULL) {
  if(length(dat[[strat]]) == 0) {
    dat$strat <- 1
  }
  strats <- unique(dat[[strat]])
  tmp <- c()
  include_ci <- FALSE
  if (!is.null(ci)) {
    if(length(ci) == 2 && (round(ci[1],3) != round((1-ci[2]),3))) {
      stop("Sorry, only symmetric confidence intervals can be computed. Please adjust the ci argument.")
    }
    include_ci <- TRUE
    if(length(ci) == 2) { # when specified as c(0.05, 0.95)
      ci <- diff(ci)
    }
  } else {
    ci <- 0.9
  }
  for (i in seq(strats)) {
    if(rtte_conditional) {
      tmp1 <- dat[dat[[strat]] == strats[i],]
    } else {
      tmp1 <- dat
      tmp1[[strat]] <- as.num(tmp1[[strat]])
      idx1 <- tmp1[[strat]] < as.num(strats[i])
      if(sum(idx1) > 0) tmp1[idx1,]$dv <- 0
      idx2 <- tmp1[[strat]] <= as.num(strats[i])
      if(sum(idx2) > 0) tmp1 <- tmp1[idx2,]
    }
    if(length(grep( "rtte", strats[1])) > 0  & i > 1) {
      for (j in 1:i) {
        tmp_j <- dat[dat[[strat]] == strats[j] & dat$dv == 0,]
        tmp1 <- rbind(tmp1, tmp_j)
      }
    }
    km_fit <- survival::survfit(survival::Surv(time = time, dv != 0) ~ 1, data = tmp1, conf.int = ci)
    km_dat <- data.frame(time = km_fit$time, surv = km_fit$surv, strat = strats[i])
    if(include_ci) {
      km_dat <- data.frame(cbind(km_dat, lower=km_fit$lower, upper=km_fit$upper))
    }
    if(reverse_prob) {
      km_dat$surv <- 1 - km_dat$surv
      if(!is.null(ci)) {
        km_dat$upper <- 1 - km_dat$upper
        km_dat$lower <- 1 - km_dat$lower
      }
    }
    first_rec <- data.frame(time=0, surv=1-as.numeric(reverse_prob), strat = strats[i])
    if(include_ci) {
      first_rec <- data.frame(cbind(first_rec, lower=1-as.numeric(reverse_prob), upper=1-as.numeric(reverse_prob)))
    }
    tmp <- rbind(tmp, rbind(first_rec, km_dat))
  }
  tmp$qmed <- tmp$surv
  tmp
}