R/outbreak_step.R

In dachuwu/DTQbp: A package for modeling case detection, contact tracing and quarantine with branching processes

#' 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 dispersion parameter for isolated cases (must be >0)
#' @param disp.com dispersion parameter for non-isolated cases (must be >0)
#' @param r0isolated reproduction number for isolated cases (must be >0)
#' @param r0symp 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 proportion of infectious contacts ascertained by contact tracing (must be 0<=x<=1)
#' @param k 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 fifelse
#' @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,
                          r0symp = NULL, disp.com = NULL, disp.iso = NULL, r0isolated = NULL,
                          prop.asym = NULL, relR.asym = NULL,
                          prop.ascertain = NULL,
                          quarant.days = NULL, quarant.retro.days = NULL,
                          incfn = NULL, delayfn = NULL, inf_fn = 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
  # + record every cases' sampled R0

  actid <- which(!(case_data$isolated))
  new_cases <- purrr::map2_dbl(
    data.table::fifelse(vect_isTRUE(case_data$isolated), disp.iso, disp.com),
    data.table::fifelse(vect_isTRUE(case_data$isolated), r0isolated, r0symp) * data.table::fifelse(vect_isTRUE(case_data$asym), relR.asym, 1),
    ~ rnbinom(1, size = .x, mu = .y))

  case_data$R0[actid] <- new_cases[actid]

  # Select cases that have generated any new cases
  new_case_data <- case_data[new_cases > 0,]
  new_cases <- new_cases[new_cases > 0]

  # The total new cases generated
  total_new_cases <- 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
    return(case_data)
  }

  # 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_cases, new_case_data$onset,
                                  function(x, y) {
                                    inf_fn(rep(y, x))
                                  })),
    # records the infector of each new person
    infector = rep(new_case_data$caseid, new_cases),

    # records when infector was isolated
    infector_iso_time = rep(new_case_data$isolated_time,new_cases),

    infector_quart_time = rep(new_case_data$quart_time, new_cases),

    infector_quart_end = rep(new_case_data$quart_end, new_cases),

    # records if infector asymptomatic
    infector_asym = rep(new_case_data$asym,new_cases),

    # 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,

    # no. of generation
    generation = rep(new_case_data$generation+1, new_cases),
    cluster = rep(new_case_data$cluster, new_cases),

    R0 = NA,
    Re = NA,
    quart_time = NA_real_,
    quart_end = NA_real_,
    isolated_time = NA_real_
  )




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

  # Set new case ids for new people
  total_new_cases <- NROW(prob_samples)
  prob_samples$caseid <- (nrow(case_data) + 1):(nrow(case_data) + total_new_cases)
  prob_samples[, cluster_obs :=  data.table::fifelse(asym, NA_integer_,
                                                     data.table::fifelse(infector_asym, caseid, cluster))]


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

  # If you are missed, you're never quarantined
  ids <- which(prob_samples$missed)
  prob_samples[ids, `:=`(
    isolated_time = data.table::fifelse(asym, Inf,
                                        # If you are not asymptomatic, but you are missed,
                                        # you are isolated at your symptom onset
                                        onset + delayfn(.N)),
    quart_time = Inf,
    quart_end = Inf
  )]


  if(quarant.days==0){

    ids <- which(!prob_samples$missed)
    prob_samples[ids, `:=`(
      isolated_time = data.table::fifelse(onset>=(infector_iso_time-1) & onset<=(infector_iso_time+1) ,
                                          vect_max(onset, infector_iso_time), onset + delayfn(.N)),
      quart_time = Inf,
      quart_end = Inf
    )]

  } else {

    ids <- which(!prob_samples$missed)

    prob_samples[ids, `:=`(
      quart_time = infector_iso_time,
      quart_end = infector_iso_time + quarant.days
    )]

    prob_samples[ids, `:=`(
      isolated_time = data.table::fifelse(onset>=(quart_time-quarant.retro.days) & onset<=quart_end,
                                          quart_time, onset + delayfn(.N))
    )]



  }


  # Chop out unneeded sample columns
  prob_samples[, c("incubfn_sample", "infector_iso_time",
                   "infector_quart_time","infector_quart_end", "infector_asym") := NULL]
  # Set new case ids for new people
  prob_samples$caseid <- (nrow(case_data) + 1):(nrow(case_data) + total_new_cases)

  # calculate individual Rj
  tmp <- table(prob_samples$infector)
  id2re <- structure(as.vector(tmp), names=names(tmp))
  case_data$Re[actid] <- id2re[as.character(case_data$caseid[actid])]



  ## Estimate the effective r0
  #effective_r0 <- total_new_cases / 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
  return(case_data)
}