R/outbreak_step.R

In covid19risk/impact-sim: Feasibility of controlling 2019-nCoV outbreaks by isolation of cases and contacts

#' 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 from the onset-to-hospitalisation delay; generated by dist_setup
#' @param prop.ascertain numeric proportion of infectious contacts ascertained by contact tracing (must be 0<=x<=1)
#' @param k numeric skew parameter for sampling the serial interval from the incubation period
#' @param quarantine logical whether quarantine is in effect, if TRUE then traced contacts are isolated before symptom onset
#'
#' @importFrom data.table data.table rbindlist
#' @importFrom purrr map2 map2_dbl map_lgl rbernoulli
#'
#' @return
#' @export
#'
#' @examples
#'
#'\dontrun{
#' # incubation period sampling function
#' incfn <- dist_setup(dist_shape = 2.322737,dist_scale = 6.492272)
#' # delay distribution sampling function
#' delayfn <- dist_setup(delay_shape, delay_scale)
#' # generate initial cases
#' case_data <- outbreak_setup(num.initial.cases = 5,incfn,delayfn,k=1.95,prop.asym=0)
#' # generate next generation of cases
#' case_data <- outbreak_step(case_data,1,0.16,0,2.5,0,incfn,delayfn,0,1.95,FALSE)
#'}
outbreak_step <- function(case_data = NULL, disp.iso = NULL, disp.com = NULL, r0isolated = NULL, r0community = NULL,
                          prop.asym = NULL, incfn = NULL, delayfn = NULL, prop.ascertain = NULL, k = NULL, quarantine = NULL, flow = NULL, traceAcc = NULL) {



  # 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)
  }

  # For each case in case_data, draw new_cases from a negative binomial distribution
  # with an R0 and dispersion dependent on if isolated=TRUE
  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), k)
                                    })),
    # 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 if infector asymptomatic
    infector_asym = unlist(purrr::map2(new_case_data$asym, new_case_data$new_cases,
                                       function(x, y) {
                                         rep(x, y)
                                         })),
    #sample to pick which groups new people are in
    infector_group = unlist(purrr::map2(new_case_data$group, new_case_data$new_cases,
                                        function(x, y) {
                                          rep(x, y)
                                        })),

    group = unlist(purrr::map2(new_case_data$group, new_case_data$new_cases,
                                        function(x, y) {
                                          sample(seq_len(nrow(flow)), y, prob=flow[x,], replace = TRUE)
                                        })),

    # 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
    # todo: modify probability based on groups
    #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[, prob_trace := traceAcc[cbind(infector_group, group)]]

  prob_samples[, missed := unlist(purrr::map(prob_samples$prob_trace,
                                              function(x) {
                                                purrr::rbernoulli(n = 1, p = 1-x)
                                              })),]



  prob_samples <- prob_samples[exposure < infector_iso_time][, # filter out new cases prevented by isolation
                                             `:=`(# onset of new case is exposure + incubation period sample
                                               onset = exposure + incubfn_sample)]


  # cases whose parents are asymptomatic are automatically missed
  prob_samples$missed[vect_isTRUE(prob_samples$infector_asym)] <- TRUE

  # If you are asymptomatic, your isolation time is Inf
  prob_samples[, isolated_time := ifelse(vect_isTRUE(asym), Inf,
                                        # If you are not asymptomatic, but you are missed,
                                        # you are isolated at your symptom onset
                                        ifelse(vect_isTRUE(missed), onset + delayfn(1),
                                               # If you are not asymptomatic and you are traced,
                                               # you are isolated at max(onset,infector isolation time) # max(onset,infector_iso_time)
                                               ifelse(!vect_isTRUE(rep(quarantine, total_new_cases)),
                                                      vect_min(onset + delayfn(1), vect_max(onset, infector_iso_time)),
                                                      infector_iso_time)))]


  # Chop out unneeded sample columns
  prob_samples[, c("incubfn_sample", "infector_iso_time", "infector_asym", "prob_trace") := 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)
}