R/outbreak_step.R

In emmalouisedavis/TTI-branching-process: TTI COVID-19 branching process model

#' Move forward one generation in the branching process
#'
#' @author Joel Hellewell
#'
#' @param case_data data.table of cases in outbreak so far; initially generated by outbreak_setup
#' @param disp.iso numeric dispersion parameter for isolated cases (must be >0)
#' @param disp.com numeric dispersion parameter for non-isolated cases (must be >0)
#' @param r0isolated numeric reproduction number for isolated cases (must be >0)
#' @param r0community numeric reproduction number for non-isolated cases (must be >0)
#' @param incfn function samples from incubation period; generated by dist_setup
#' @param delayfn function samples adherence and delay to isolation
#' @param prop.ascertain numeric proportion of infectious contacts ascertained by contact tracing (must be 0<=x<=1)
#' @param quarantine logical whether quarantine is in effect, if TRUE then traced contacts are isolated before symptom onset
#' @param min_quar_delay The minimum delay between a case being identified and their contacts being isolated (only applies when quarentine set to TRUE)
#' @param max_quar_delay The maximum delay between a case being identified and their contacts being isolated (only applies when quarentine set to TRUE)
#' @param prop.asym Proportion of asymptomatics.
#' @param inf_rate Rate parameter for the gamma distribution of serial intervals around the symptom onset distribution.
#' @param inf_shape The shape for the gamma distribution of serial intervals around the symptom onset distribution.
#' @param inf_shift Shift the gamma distribution of serial intervals around the symptom onset distribution back by this much (i.e. transmission can ocur this much before symptom onset).
#' @param test_delay How long does it take for tests to be administered and results returned.
#' @param sensitivity Test sensitivity.
#' @param precaution After a negative test result, keep people in quarantine for this long as a precautionary measure.
#' @param self_report Probability that someone that is not tracked will self report (111 for example) after symptoms.
#' @param testing Logical to determine whether testing is used.
#' @param iso_adhere Probability an individual will follow advice from TTI to isolate.
#' @param test_asym Logical to determine whether asymptomatics are tested.
#' @importFrom data.table data.table rbindlist
#' @importFrom purrr map2 map2_dbl map_lgl rbernoulli
#'
#' @return The new cases data.table.
#' @export
#'
#'
outbreak_step <- function(case_data, disp.iso,
                          disp.com, r0isolated,
                          r0community, prop.asym,
                          incfn, delayfn,
                          inf_rate, inf_shape,
                          inf_shift, prop.ascertain,
                          min_quar_delay, max_quar_delay,
                          test_delay, sensitivity,
                          precaution, self_report,
                          quarantine, testing,
                          test_asym, iso_adhere) {

  # Column names used in nonstandard eval.
  test_result <- isolated_end <- infector_iso_end <- delays <- NULL
  delays_traced <- test <- time_to_test <- test_result <- isolated_end <- NULL

  # For each case in case_data, draw new_cases from a negative binomial distribution
  # with an R0 and dispersion dependent on if isolated=TRUE
  new_cases <- case_data[, new_cases := purrr::map2_dbl(
    ifelse(vect_isTRUE(isolated), disp.iso, disp.com),
    ifelse(vect_isTRUE(isolated),
           r0isolated,
           r0community),
    ~ rnbinom(1, size = .x, mu = .y))
    ]

  # Select cases that have generated any new cases
  new_case_data <- case_data[new_cases > 0]
  # The total new cases generated
  total_new_cases <- case_data[, sum(new_cases), ]

  # If no new cases drawn, outbreak is over so return case_data
  if (total_new_cases == 0) {
    # If everyone is isolated it means that either control has worked or everyone has had a chance to infect but didn't
    case_data$isolated <- TRUE

    effective_r0 <- 0
    cases_in_gen <- 0
    out <- list(case_data, effective_r0, cases_in_gen)
    names(out) <- c("cases", "effective_r0", "cases_in_gen")

    return(out)
  }

  # Compile a data.table for all new cases, new_cases is the amount of people that each infector has infected
  inc_samples <- incfn(total_new_cases)

  prob_samples <- data.table(
    # time when new cases were exposed, a draw from serial interval based on infector's onset
    exposure = unlist(purrr::map2(new_case_data$new_cases, new_case_data$onset,
                                  function(x, y) {
                                    inf_fn(rep(y, x), inf_shape, inf_rate, inf_shift)
                                    })),
    # records the infector of each new person
    infector = unlist(purrr::map2(new_case_data$caseid, new_case_data$new_cases,
                                  function(x, y) {
                                    rep(as.integer(x), as.integer(y))
                                    })),
    # records when infector was isolated
    infector_iso_time = unlist(purrr::map2(new_case_data$isolated_time, new_case_data$new_cases,
                                           function(x, y) {
                                             rep(x, as.integer(y))
                                             })),
    # records when infector came out of isolation
    infector_iso_end = unlist(purrr::map2(new_case_data$isolated_end, new_case_data$new_cases,
                                           function(x, y) {
                                             rep(x, as.integer(y))
                                           })),
    # records if infector asymptomatic
    infector_asym = unlist(purrr::map2(new_case_data$asym, new_case_data$new_cases,
                                       function(x, y) {
                                         rep(x, y)
                                         })),
    # records if infector has tested positive
    infector_pos = unlist(purrr::map2(new_case_data$test_result, new_case_data$new_cases,
                                       function(x, y) {
                                         rep(x, y)
                                       })),
    # records time infector was tested
    infector_testtime = unlist(purrr::map2(new_case_data$test_time, new_case_data$new_cases,
                                      function(x, y) {
                                        rep(x, y)
                                      })),
    # records infector onset time
    infector_onset = unlist(purrr::map2(new_case_data$onset, new_case_data$new_cases,
                                      function(x, y) {
                                        rep(x, y)
                                      })),
    # records if infector was missed
    infector_missed = unlist(purrr::map2(new_case_data$missed, new_case_data$new_cases,
                                      function(x, y) {
                                        rep(x, y)
                                      })),
    # records if infector refused to isolate
    infector_refuse = unlist(purrr::map2(new_case_data$refuse, new_case_data$new_cases,
                                         function(x, y) {
                                           rep(x, y)
                                         })),
    # records when infector was exposed
    infector_exp = unlist(purrr::map2(new_case_data$exposure, new_case_data$new_cases,
                                      function(x, y) {
                                        rep(x, y)
                                      })),
    # draws a sample to see if this person is asymptomatic
    asym = purrr::rbernoulli(n = total_new_cases, p = prop.asym),
    # draws a sample to see if this person is traced
    missed = purrr::rbernoulli(n = total_new_cases, p = 1 - prop.ascertain),
    # sample from the incubation period for each new person
    incubfn_sample = inc_samples,
    isolated = FALSE,
    new_cases = NA
  )

  prob_samples$exposure <- pmax(prob_samples$exposure, prob_samples$infector_exp +1)
  prob_samples <- prob_samples[exposure < infector_iso_time | exposure > infector_iso_end | (infector_refuse == TRUE & infector_missed == FALSE)][, # filter out new cases prevented by isolation
                                             `:=`(# onset of new case is exposure + incubation period sample
                                               onset = exposure + incubfn_sample)]

  # cases whose parents have been missed are automatically missed
  prob_samples$missed[vect_isTRUE(prob_samples$infector_missed)] <- TRUE
  # cases who were infected more than 3 days before infector's symptoms can't be traced
  prob_samples$missed[vect_isTRUE(prob_samples$exposure < prob_samples$infector_onset - 3)] <- TRUE

  ##UK tracing edit 27/05/2020:
  # cases whose parents tested negative are automatically missed
  prob_samples$missed[vect_isTRUE(!prob_samples$infector_pos)] <- TRUE

  #had to put these outside as the vectorisation wasn't working inside the next line down where isolate_time is calculated, was using the same sampled value for all columns
  prob_samples[, delays := delayfn(nrow(prob_samples))] #delays of symptomatic individuals to isolation
  prob_samples[, delays_traced := runif(nrow(prob_samples), min_quar_delay, max_quar_delay)] #delays of those who are traced

  prob_samples[, refuse := rbinom(nrow(prob_samples),1,1-iso_adhere)] #who will refuse to isolated even if traced/detected

  prob_samples[, missed := ifelse(vect_isTRUE(exposure > infector_iso_end), #Individuals exposed after infector released from quarantine due to FN test aren't traced
                                  TRUE,missed)]

  # If you are asymptomatic, your isolation time is Inf
  prob_samples[, isolated_time := ifelse(vect_isTRUE(missed),
                                         # If you are not tracked (are missed)
                                         ifelse(vect_isTRUE(asym),
                                                # If you a are missed asymptotic you never isolate
                                                Inf,
                                                # If you are not asymptomatic you isolate after onset of symptoms plus a delay
                                                onset + delays), # delays = Inf if non-adherent
                                         # If you are tracked (are not missed)
                                         ifelse(vect_isTRUE(rep(quarantine, total_new_cases)),
                                                # With quarantine, isolate with some delay after your infector was identified.
                                                ifelse(vect_isTRUE(asym),
                                                       infector_testtime + delays_traced, #infector_test_time + delays_traced (same next line)
                                                       vect_min(onset+delays,infector_testtime + delays_traced)), #minimum of isolation due to symptoms and isolation due to tracing
                                                # Without quarantine:
                                                # onset < infector_iso < onset+delay  -> isolate when infector is identified and isolates
                                                # onset < onset+delay  < infector_iso -> isolate after symptoms and a delay
                                                # infector_iso < onset < onset+delay  -> isolate as soon as symptoms onset.
                                                vect_min(onset + delayfn(1), vect_max(onset, infector_iso_time))))]


  missedSympt <- nrow(prob_samples[vect_isTRUE(missed) & vect_isTRUE(!asym),]) # Symptomatic individuals who are missed
  prob_samples[vect_isTRUE(missed) & vect_isTRUE(!asym), missed:=purrr::rbernoulli(missedSympt,p=1-self_report)] # Report themselves to contact tracing with prob self_report

  if(test_asym == FALSE){
    prob_samples$missed[vect_isTRUE(prob_samples$asym)] <- TRUE #if asymptomatic then missed and therefore not tested. Important that this is AFTER isolated time calculations as will still be asked to isolate if traced.
  }

  if(testing==TRUE) {
      prob_samples[, test := ifelse(vect_isTRUE(missed), # & vect_isTRUE(!asym), # If not-traced:
                                    FALSE, # not tested
                                    TRUE)] # otherwise tested
      if(test_asym==TRUE){
        prob_samples[, time_to_test := ifelse(vect_isTRUE(test), # If tested:
                                            test_delay, # delay from isolation to test results (currently a constant delay but could be expanded)
                                            Inf)]
      }
      if(test_asym==FALSE){
        prob_samples[, time_to_test := ifelse(vect_isTRUE(test), # If tested:
                                            pmax(0,onset-isolated_time)+test_delay, # delay from onset (or isolation if isolation after onset) to test results, assuming only test when symptoms appear (currently a constant delay but could be expanded)
                                            Inf)]
      }

      prob_samples[, test_time := ifelse(vect_isTRUE(test), # If tested:
                                          isolated_time+time_to_test, # actual time test result received
                                          Inf)]

      prob_samples[, test_result := ifelse(vect_isTRUE(test), # If tested
                                            as.logical(rbinom(length(which(prob_samples$test==T)),1,sensitivity)), # =TRUE if positive, =FALSE if false negative
                                            NA)] # not tested

      prob_samples[, isolated_end := ifelse(vect_isTRUE(test), # If tested
                                            ifelse(vect_isTRUE(test_result), # If positive
                                              Inf, # Stay in isolation long enough to not transmit
                                              vect_max(isolated_time+time_to_test,isolated_time+precaution)), #onset + time_to_test # If test is negative
                                            # Leave isolation with some precautionary delay (0-7 days)
                                            Inf)]

      # prob_samples[vect_isTRUE(!prob_samples$infector_pos) & vect_isTRUE(!missed), #if you were traced but your infector didn't test positive
      #              isolated_end := ifelse((infector_iso_time + test_delay)<=isolated_time, #if their test came back before you were traced and isolated
      #                                     isolated_time+precaution, #then you stay in isolation for time = precaution
      #                                     ifelse(vect_isTRUE(test_result), #if you were isolated before their test then you're also tested
      #                                            isolated_end, #if you tested positive then no change to end of isolation (i.e. =Inf)
      #                                            vect_max(isolated_time+time_to_test,isolated_time+precaution)))] #if you tested negative then you are released after your test result, providing you've been in isolation for at least precaution days
      #
      # prob_samples[vect_isTRUE(prob_samples$infector_pos==FALSE), #if your infector tested negative
      #              missed := ifelse((infector_iso_time + test_delay)<=isolated_time, #if their test came back before you were traced
      #                               TRUE, #your contacts aren't traced
      #                               ifelse(vect_isTRUE(test_result), #otherwise, you're tested
      #                                      missed, #positive: no change to previous allocation
      #                                      TRUE))] #negative: your contacts also aren't traced
  }

  if(testing == FALSE) {
    prob_samples$test <- FALSE
    prob_samples$time_to_test <- NA
    prob_samples$test_time <- NA
    prob_samples$test_result <- NA
    prob_samples$isolated_end <- NA
    prob_samples$missed[vect_isTRUE(prob_samples$infector_asym)] <- TRUE
  }

  #make sure no one has isolation start time after their isolation end time
  prob_samples[, isolated_end := ifelse(vect_isTRUE(isolated_time > isolated_end),
                                        isolated_time,
                                        isolated_end)]

  #prob_samples[vect_isTRUE(isolated_time>=isolated_end),isolated_time := Inf]
  #prob_samples[vect_isTRUE(isolated_time>=isolated_end),isolated_end := Inf]

  # Chop out unneeded sample columns
  prob_samples[, c("incubfn_sample", "infector_iso_time", "infector_asym","infector_pos",
                   "infector_testtime","infector_exp","infector_iso_end","delays","delays_traced","test",
                   'time_to_test',"infector_missed","infector_refuse","infector_onset") := NULL]
  # Set new case ids for new people
  prob_samples$caseid <- (nrow(case_data) + 1):(nrow(case_data) + nrow(prob_samples))

  ## Number of new cases
  cases_in_gen <- nrow(prob_samples)

  ## Estimate the effective r0
  effective_r0 <- nrow(prob_samples) / nrow(case_data[vect_isTRUE(!case_data$isolated)])

  # Everyone in case_data so far has had their chance to infect and are therefore considered isolated
  case_data$isolated <- TRUE

  # bind original cases + new secondary cases
  case_data <- data.table::rbindlist(list(case_data, prob_samples),
                                     use.names = TRUE)

  # Return
  out <- list(case_data, effective_r0, cases_in_gen)
  names(out) <- c("cases", "effective_r0", "cases_in_gen")

  return(out)
}

# A vectorised version of isTRUE
vect_isTRUE <- function(x) {
  purrr::map_lgl(x, isTRUE)
}


vect_max <- function(x, y) {
  purrr::map2_dbl(x, y, max)
}

vect_min <- function(x, y) {
  purrr::map2_dbl(x, y, min)
}