R/pf_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 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 r0symp 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 quarant.days quarantine days
#' @param quarant.retro.days batracking of symptom onset
#'
#' @importFrom data.table data.table rbindlist fifelse
#' @importFrom purrr map2 map2_dbl map_lgl rbernoulli
#'
#' @return
#' @export
#'
pf_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] <- case_data$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 generation 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, as.integer(new_cases)),

    # records of infector
    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),

    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_
  )

  # filter out new cases prevented by isolation  or quarantine
  prob_samples <- prob_samples[(exposure < infector_iso_time) & (exposure < infector_quart_time | exposure > infector_quart_end)]

  # Set new case ids for new people
  total_new_cases <- NROW(prob_samples)
  exist_cases <- NROW(case_data)
  prob_samples$caseid <- (exist_cases + 1):(exist_cases + total_new_cases)
  prob_samples[, `:=`(
    onset = exposure + incubfn_sample,
    cluster_obs =  data.table::fifelse(asym,
                                       NA_integer_,
                                       data.table::fifelse(infector_iso_time==Inf, 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(length(ids))),
   ## cluster_obs = data.table::fifelse(is.na(cluster_obs), NA_integer_, caseid),  ####
    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(length(ids))),
      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(length(ids)))
    )]

  }

  # ids <- which(prob_samples$isolated_time == Inf)
  # prob_samples[ids, `:=`( cluster_obs = NA_integer_ )]


  # Chop out unneeded sample columns
  prob_samples[, c("exposure", "incubfn_sample", "infector_iso_time",
                   "infector_quart_time","infector_quart_end", "infector_asym") := NULL]


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

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