R/branching_model.R

In RussellXu/rtestimate: Estimation Of Reproductive Number

#' Branching process model. Random simulation according to the required parameters
#'
#' @param num.initial.cases the number of initial cases
#' @param prop.ascertain the proportion of transmission that occurred before symptom onset
#' @param cap_max_days maximum times
#' @param cap_cases maximum cases
#' @param r0isolated basic reproduction number during isolation period
#' @param r0community basic reproduction number during non-isolation period
#' @param disp.iso dispersion parameter during isolation period
#' @param disp.com dispersion parameter during non-isloation period
#' @param delay_shape the serial interval distribution shape parameter of the delay from symptom onset to isolation
#' @param delay_scale the serial interval distribution scale parameter of the delay from symptom onset to isolation
#' @param prop.asym the proportion of asymptomatic
#' @param quarantine logical value
#'
#' @importFrom data.table rbindlist
#' @import magrittr
#'
#' @family branching process
#'
#' @return dayly_cases

branching_model <- function(num.initial.cases = 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, prop.asym = NULL,
                           quarantine = NULL) {

  # Set up functions to sample from distributions
  # incubation period sampling function
  incfn <- dist_setup(dist_shape = 2.322737,
                      dist_scale = 6.492272)
  # incfn <- dist_setup(dist_shape = 3.303525,dist_scale = 6.68849) # incubation function for ECDC run
  # onset to isolation delay sampling function
  delayfn <- dist_setup(delay_shape,
                        delay_scale)

  # Set initial values for loop indices
  total.cases <- num.initial.cases
  latest.onset <- 0
  extinct <- FALSE

  # Initial setup
  case_data <- branching_setup(num.initial.cases = num.initial.cases,
                              incfn = incfn,
                              prop.asym = prop.asym,
                              delayfn = delayfn,
                              k = k)

  # Preallocate
  effective_r0_vect <- c()
  cases_in_gen_vect <- c()


  # Model loop
  while (latest.onset < cap_max_days & total.cases < cap_cases & !extinct) {

    out <- branching_step(case_data = case_data,
                         disp.iso = disp.iso,
                         disp.com = disp.com,
                         r0isolated = r0isolated,
                         r0community = r0community,
                         incfn = incfn,
                         delayfn = delayfn,
                         prop.ascertain = prop.ascertain,
                         k = k,
                         quarantine = quarantine,
                         prop.asym = prop.asym)


    case_data <- out[[1]]
    effective_r0_vect <- c(effective_r0_vect, out[[2]])


    cases_in_gen_vect <- c(cases_in_gen_vect, out[[3]])
    total.cases <- nrow(case_data)
    latest.onset <- max(case_data$onset)
    extinct <- all(case_data$isolated)
  }

  r0 <- Re_estimate(case_data = case_data, cap_max_days = cap_max_days)


  # Prepare output, group into days
  dayly_cases <- case_data[, day := floor(onset)
  ][, .(dayly_cases = .N), by = day]
  # maximum outbreak day
  max_day <- floor(cap_max_days)
  # days with 0 cases in 0:max_day
  missing_days <- (0:max_day)[!(0:max_day %in% dayly_cases$day)]

  # add in missing days if any are missing
  if (length(missing_days > 0)) {
    dayly_cases <- data.table::rbindlist(list(dayly_cases,
                                               data.table(day = missing_days,
                                                          dayly_cases = 0)))
  }
  # order and sum up
  dayly_cases <- dayly_cases[order(day)
  ][, cumulative := cumsum(dayly_cases)]
  # cut at max_day
  dayly_cases <- dayly_cases[day <= max_day]

  # Add effective R0
  dayly_cases <- dayly_cases[, `:=`(cases_per_gen = list(cases_in_gen_vect))]

  dayly_cases <- dayly_cases[, effective_r0 := r0]
  # return
  return(dayly_cases)
}


#' Branching Rt estimation function.The core algorithm in calculating Rt values.
#'
#' @param case_data
#' @param cap_max_days
#'
#' @author Xu Lingfeng, \email{1760019613@qq.com}
#'
#' @family branching process
#'
#' @return time-varying Rt list

Re_estimate <- function(case_data, cap_max_days) {

  r0_case_data <- case_data
  r0_case_data <- r0_case_data[, exposure := floor(exposure)]
  r0_case_data <- r0_case_data[, isolated_time := floor(isolated_time)]

  r0_case_data <- data.table(exposure = case_data$exposure,
                             caseid = r0_case_data$caseid,
                             infector = r0_case_data$infector,
                             isolated_time = case_data$isolated_time)

  # 计算Rt：以原始infector为分母，他们在某段时间内感染的人数为分子。
  # if(i == 1){
  #
  # }
  # y <- r0_case_data[exposure <=i - 1]
  # x <- r0_case_data[exposure <=i]
  #
  # if(identical(x, y)){
  #   r0[i] <- r0[i-1]
  # } else {
  #   temp_data <- r0_case_data[exposure <= i & isolated_time > i]
  #   infector <- nrow(temp_data[!infector %in% caseid])
  #   infectee <- nrow(temp_data)
  #   r0[i] <- infectee/infector
  #   }

  r0 <- c()
  for (i in 0:cap_max_days){
    # 上帝视角，以今天新产生的人，能在未来的日子里一共感染多少人计算，但只计算第二代genaration.
    infector_id <- r0_case_data[exposure ==i][,caseid]  # 这些人在整个传染过程中传染了哪些人
    new_cases_by_infector <- nrow(r0_case_data[infector %in% infector_id])
    if(length(infector_id) == 0){
      r0[i] <- 0
      next
    }
    r0[i] <- new_cases_by_infector / length(infector_id)
  }
  r0 <- append(r0, 0, after = 0)
  return(r0)

}