R/sim_SIR.R

In collinmmccabe/enss: Effective Network Size Simulation

#' sim_SIR
#'
#' sim_SIR simulates the spread of an infectious disease on a social network
#' using the Susceptible-Infected-Removed compartmental model
#'
#' @param network_el an edgelist file generated by igraph
#' @param beta a value from 0 to 1 indicating per contract probability of
#'   pathogen transmission
#' @param gamma a value from 0 to 1 indicating daily probability of
#'   recovery/removal
#' @param intxn_per_day the per capita interaction rate
#' @param days the maximum number of days for which your simulation will run
#'
#' @return vector with 5 elements, 1: number of days simulation went for, 2:
#'   final number of susceptible individuals, 3: final number of infected
#'   individuals, 4: final number of number of removed individuals, 5: maximum
#'   number of infected inviduals at any timepoint
#' @export
#'
#' @examples
sim_SIR <- function(network_el, beta, gamma, intxn_per_day, days) {

  n = network_el[1,4]
  e = network_el[1,5]
  cdata <- network_el[,1:2]

  infection_status <- c(rep(1,n))
  index_infected <- sample(1:n, 1)
  infection_status[index_infected] = 2

# only SIR
  max_infected <- 1

  day_counter <- 0
  while(day_counter <= days) {

    int_counter <- 0
    while(int_counter <= intxn_per_day*n) {

      selected_edge <- sample(1:e,1)

      if (sum(infection_status[cdata[selected_edge,1:2]]) == 3) {
        if (beta >= runif(1,0,1)) {
          infection_status[cdata[selected_edge,1:2]] = 2
        }
      }

      int_counter <- sum(int_counter,1)
    }

    # only SIR
    for (j in which(infection_status %in% 2)) {
      if (gamma >= runif(1,0,1)) {
        infection_status[j] = 3
      }
    }

    #only in SIR
    curr_infected <- sum(infection_status == 2)
    if (curr_infected > max_infected) {
      max_infected <- curr_infected
    }

    day_counter <- sum(day_counter,1)

    # 0 in SI, n in SIR
    if (sum(infection_status%%2) == n) break
  }

  # returns different info for SI and SIR
  return(c(day_counter-1,sum(infection_status == 1),sum(infection_status == 2),sum(infection_status == 3),max_infected))
}