bootstraps <- 20

split.rule <- 'logrank'

n.threads <- 16

# Cross-validation variables

ns.splits <- 0:20

ns.trees  <- c(500, 1000, 2000)

ns.imputations <- 1:3

cv.n.folds <- 3

n.data <- NA # This is of full dataset...further rows may be excluded in prep

calibration.filename <- paste0(output.filename.base, '-calibration.csv')

continuous.vars <-

  c(

    'age', 'total_chol_6mo', 'hdl_6mo', 'pulse_6mo', 'crea_6mo',

    'total_wbc_6mo', 'haemoglobin_6mo'

  )

untransformed.vars <- c('anonpatid', 'surv_time', 'imd_score', 'exclude')

source('../lib/shared.R')

require(ggrepel)

# Load the data and convert to data frame to make column-selecting code in

# prepData simpler

COHORT.full <- data.frame(fread(data.filename))

# If n.data was specified...

if(!is.na(n.data)){

  # Take a subset n.data in size

  COHORT.use <- sample.df(COHORT.full, n.data)

  rm(COHORT.full)

} else {

  # Use all the data

  COHORT.use <- COHORT.full

  rm(COHORT.full)

}

# We now need a quick null preparation of the data to get its length (some rows

# may be excluded during preparation)

COHORT.prep <-

  prepData(

    COHORT.use,

    cols.keep, discretise.settings, surv.time, surv.event,

    surv.event.yes, extra.fun = caliberExtraPrep, n.keep = n.data

  )

n.data <- nrow(COHORT.prep)

# Define indices of test set

test.set <- testSetIndices(COHORT.prep, random.seed = 78361)

# Process settings: don't touch anything!!

process.settings <-

  list(

    var       = c(untransformed.vars, continuous.vars),

    method    = rep(NA, length(untransformed.vars) + length(continuous.vars)),

    settings  = rep(NA, length(untransformed.vars) + length(continuous.vars))

  )

# If we've not already done a calibration, then do one

if(!file.exists(calibration.filename)) {

  # Create an empty data frame to aggregate stats per fold

  cv.performance <- data.frame()

  # Items to cross-validate over

  cv.vars <- expand.grid(ns.splits, ns.trees, ns.imputations)

  names(cv.vars) <- c('n.splits', 'n.trees', 'n.imputations')

  # prep the data (since we're not cross-validating on data prep this can be

  # done before the loop)

  # Prep the data

  COHORT.cv <-

    prepData(

      # Data for cross-validation excludes test set

      COHORT.use[-test.set, ],

      cols.keep,

      process.settings,

      surv.time, surv.event,

      surv.event.yes,

      extra.fun = caliberExtraPrep

    )

  # Finally, add missing flag columns, but leave the missing data intact because

  # rfsrc can do on-the-fly imputation

  COHORT.cv <- prepCoxMissing(COHORT.cv, missingReplace = NA)

  # Add on those column names we just created

  surv.predict <-

    c(surv.predict, names(COHORT.cv)[grepl('_missing', names(COHORT.cv))])

  # Run crossvalidations. No need to parallelise because rfsrc is parallelised

  for(i in 1:nrow(cv.vars)) {

    cat(

      'Calibration', i, '...\n'

    )

    # Get folds for cross-validation

    cv.folds <- cvFolds(nrow(COHORT.cv), cv.n.folds)

    cv.fold.performance <- data.frame()

    for(j in 1:cv.n.folds) {

      time.start <- handyTimer()

      # Fit model to the training set

      surv.model.fit <-

        survivalFit(

          surv.predict,

          COHORT.cv[-cv.folds[[j]],],

          model.type = 'rfsrc',

          n.trees = cv.vars$n.trees[i],

          split.rule = split.rule,

          n.threads = n.threads,

          nsplit = cv.vars$n.splits[i],

          nimpute = cv.vars$n.imputations[i],

          na.action = 'na.impute'

        )

      time.learn <- handyTimer(time.start)

      time.start <- handyTimer()

      # Get C-indices for training and validation sets

      c.index.train <-

        cIndex(

          surv.model.fit, COHORT.cv[-cv.folds[[j]],],

          na.action = 'na.impute'

        )

      c.index.val <-

        cIndex(

          surv.model.fit, COHORT.cv[cv.folds[[j]],],

          na.action = 'na.impute'

        )

      time.predict <- handyTimer(time.start)

      # Append the stats we've obtained from this fold

      cv.fold.performance <-

        rbind(

          cv.fold.performance,

          data.frame(

            calibration = i,

            cv.fold = j,

            n.trees = cv.vars$n.trees[i],

            n.splits = cv.vars$n.splits[i],

            n.imputations = cv.vars$n.imputations[i],

            c.index.train,

            c.index.val,

            time.learn,

            time.predict

          )

        )

    } # End cross-validation loop (j)

    # rbind the performance by fold

    cv.performance <-

      rbind(

        cv.performance,

        cv.fold.performance

      )

    # Save output at the end of each loop

    write.csv(cv.performance, calibration.filename)

  } # End calibration loop (i)

} else { # If we did previously calibrate, load it

  cv.performance <- read.csv(calibration.filename)

}

# Find the best calibration...

# First, average performance across cross-validation folds

cv.performance.average <-

  aggregate(

    c.index.val ~ calibration,

    data = cv.performance,

    mean

  )

# Find the highest value

best.calibration <-

  cv.performance.average$calibration[

    which.max(cv.performance.average$c.index.val)

  ]

# And finally, find the first row of that calibration to get the n.bins values

best.calibration.row1 <-

  min(which(cv.performance$calibration == best.calibration))

# Prep the data to fit and test with

COHORT.prep <-

  prepData(

    # Data for cross-validation excludes test set

    COHORT.use,

    cols.keep,

    process.settings,

    surv.time, surv.event,

    surv.event.yes,

    extra.fun = caliberExtraPrep

  )

# Finally, add missing flag columns, but leave the missing data intact because

# rfsrc can do on-the-fly imputation

COHORT.prep <- prepCoxMissing(COHORT.prep, missingReplace = NA)

# Add on those column names we just created

surv.predict <-

  c(surv.predict, names(COHORT.prep)[grepl('_missing', names(COHORT.prep))])

#' ## Fit the final model

#' 

#' This may take some time, so we'll cache it if possible...

#+ fit_final_model

surv.model.fit <-

  survivalFit(

    surv.predict,

    COHORT.prep[-test.set,],

    model.type = 'rfsrc',

    n.trees = cv.performance[best.calibration.row1, 'n.trees'],

    split.rule = split.rule,

    n.threads = n.threads,

    nsplit = cv.performance[best.calibration.row1, 'n.splits'],

    nimpute = cv.performance[best.calibration.row1, 'n.imputations'],

    na.action = 'na.impute'

  )

surv.model.params.boot <-

  survivalFitBoot(

    surv.predict,

    COHORT.prep[-test.set,], # Training set

    COHORT.prep[test.set,],  # Test set

    model.type = 'rfsrc',

    n.threads = n.threads,

    bootstraps = bootstraps,

    n.trees = cv.performance[best.calibration.row1, 'n.trees'],

    split.rule = split.rule,

    nsplit = cv.performance[best.calibration.row1, 'n.splits'],

    nimpute = cv.performance[best.calibration.row1, 'n.imputations'],

    na.action = 'na.impute',

    filename = paste0(output.filename.base, '-boot-all.csv')

  )

# Get C-indices for training and test sets

surv.model.fit.coeffs <- bootStatsDf(surv.model.params.boot)

# Save them to the all-models comparison table

varsToTable(

  data.frame(

    model = 'rfsrc',

    imputation = FALSE,

    discretised = FALSE,

    c.index = surv.model.fit.coeffs['c.index', 'val'],

    c.index.lower = surv.model.fit.coeffs['c.index', 'lower'],

    c.index.upper = surv.model.fit.coeffs['c.index', 'upper'],

    calibration.score = surv.model.fit.coeffs['calibration.score', 'val'],

    calibration.score.lower =

      surv.model.fit.coeffs['calibration.score', 'lower'],

    calibration.score.upper =

      surv.model.fit.coeffs['calibration.score', 'upper']

  ),

  performance.file,

  index.cols = c('model', 'imputation', 'discretised')

)

write.csv(

  surv.model.fit.coeffs,

  paste0(output.filename.base, '-boot-summary.csv')

)