R/branching_step.R

In RussellXu/rtestimate: Estimation Of Reproductive Number

#' Set up the config of branching model
#'
#' @param case_data
#' @param disp.iso
#' @param disp.com
#' @param r0isolated
#' @param r0community
#' @param prop.asym
#' @param incfn
#' @param delayfn
#' @param prop.ascertain
#' @param k
#' @param quarantine
#' @return
#' @export
#'
#' @examples
#'
#' @importFrom data.table rbindlist
#' @importFrom purrr map2 map2_dbl map_lgl rbernoulli
#'
#' @family branching process
branching_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) {

  # 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)
                                       })),
    # 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 <- 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") := 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)
}


#' Calling the branching simulation
#'
#' @param n.sim
#' @param prop.ascertain
#' @param cap_max_days
#' @param cap_cases
#' @param r0isolated
#' @param r0community
#' @param disp.iso
#' @param disp.com
#' @param k
#' @param delay_shape
#' @param delay_scale
#' @param num.initial.cases
#' @param prop.asym
#' @param quarantine
#'
#' @importFrom purrr safely
#' @return
#' @export
#'
#' @examples
#' @family branching process
branching_simulation <- function(n.sim = NULL, prop.ascertain = NULL, cap_max_days = NULL, cap_cases = NULL,
                                 r0isolated = NULL, r0community = NULL, disp.iso = NULL, disp.com = NULL, k = NULL,
                                 delay_shape = NULL, delay_scale = NULL, num.initial.cases = NULL,
                                 prop.asym = NULL, quarantine = NULL) {

  # Run n.sim number of model runs and put them all together in a big data.frame
  res <- purrr::map(.x = 1:n.sim, ~ branching_model(num.initial.cases = num.initial.cases,
                                                    prop.ascertain = prop.ascertain,
                                                    cap_max_days = cap_max_days,
                                                    cap_cases = cap_cases,
                                                    r0isolated = r0isolated,
                                                    r0community = r0community,
                                                    disp.iso = disp.iso,
                                                    disp.com = disp.com,
                                                    delay_shape = delay_shape,
                                                    delay_scale = delay_scale,
                                                    k = k,
                                                    prop.asym = prop.asym,
                                                    quarantine = quarantine))


  # bind output together and add simulation index
  res <- data.table::rbindlist(res)
  res[, sim := rep(1:n.sim, rep(floor(cap_max_days) + 1, n.sim)), ]
  return(res)
}


#' Distribution setup
#'
#' @param dist_shape
#' @param dist_scale
#' @importFrom purrr partial
#'
#' @return
#' @family branching process
dist_setup <- function(dist_shape = NULL, dist_scale = NULL) {
  out <- purrr::partial(rweibull,
                        shape = dist_shape,
                        scale = dist_scale)
  return(out)
}


#' Skew-Normal Distribution set up
#'
#' @param inc_samp
#' @param k
#' @importFrom sn rsn
#' @return
#'
#' @family branching process
#'
inf_fn <- function(inc_samp = NULL, k = NULL) {

  out <- sn::rsn(n = length(inc_samp),
                 xi = inc_samp,
                 omega = 2,
                 alpha = k)

  out <- ifelse(out < 1, 1, out)

  return(out)
}


#' The proportion of epidemic extinction
#'
#' @param outbreak_df_week
#' @param cap_cases
#'
#' @return
#' @export
#'
#' @examples
#'
#' @family branching process
#'
extinct_prob <- function(outbreak_df_week  = NULL, cap_cases  = NULL) {

  n_sim <- max(outbreak_df_week$sim)

  out <- outbreak_df_week %>%
    # new variable extinct = 1 if cases in weeks 10-12 all 0, 0 if not
    detect_extinct(cap_cases) %>%
    # number of runs where extinct = TRUE / number of runs
    .$extinct %>%
    sum(.) / n_sim

  return(out)
}



#' Detect extinction of process of branching simulation
#'
#' @param outbreak_df_week
#' @param cap_cases
#' @importFrom dplyr group_by filter summarise ungroup
#' @return
#' @export
#'
#' @examples
#' @family branching process
detect_extinct <- function(outbreak_df_week  = NULL, cap_cases  = NULL) {

  outbreak_df_week %>%
    dplyr::group_by(sim) %>% # group by simulation run
    dplyr::filter(week %in% 12:16) %>%
    dplyr::summarise(extinct =
                       ifelse(all(weekly_cases == 0 &
                                    cumulative < cap_cases),
                              1, 0)) %>%
    dplyr::ungroup()

}


#' Branching method setup
#'
#' @param num.initial.cases
#' @param incfn
#' @param delayfn
#' @param k
#' @param prop.asym
#' @importFrom data.table data.table
#'
#' @return
#' @export
#'
#' @examples
#'
#'@family branching process
branching_setup <- function(num.initial.cases, incfn, delayfn, k, prop.asym) {
  # Set up table of initial cases
  inc_samples <- incfn(num.initial.cases)

  case_data <- data.table(exposure = rep(0, num.initial.cases), # Exposure time of 0 for all initial cases
                          asym = purrr::rbernoulli(num.initial.cases, prop.asym),
                          caseid = 1:(num.initial.cases), # set case id
                          infector = 0,
                          missed = TRUE,
                          onset = inc_samples,
                          new_cases = NA)

  # set isolation time for cluster to minimum time of onset of symptoms + draw from delay distribution
  case_data <- case_data[, isolated_time := onset + delayfn(1)
  ][, isolated := FALSE]

  case_data$isolated_time[case_data$asym] <- Inf

  # return
  return(case_data)
}