simnetR: Epidemic Simulations on Networks

#' Initial Rates of Individuals in the Network
#'
#' \code{init_rates_sis_matrix} returns an array of the initial rates associated with each
#' individual in the network.
#'
#' @param N Number of vertices / nodes in the network.
#' @param startnodes Initially infectious individuals.
#' @param gam recovery rate
#' @param tau per link infection rate
#' @param graph the adjacency matrix representation of the network.
#' @param status the initial infectious status of individuals.
#' @keywords internal
#' @return Integer array of size \code{N} representing the rates applicable to
#'   individuals
#' @export
init_rates_sis_matrix <- function(graph, startnodes, gam, tau, status) {
  N <- nrow(graph)
  rate <- integer(N)
  rate[startnodes] <- gam

  for (i in startnodes) {
    s_contacts <- sus_contacts_matrix(graph[, i], status)
    rate[s_contacts] <- rate[s_contacts] + tau
  }

  return(rate)
}



#' Initial Rates of Individuals in the Network
#'
#' \code{init_rates_sis_list} returns an array of the initial rates associated with each
#' individual in the network.
#'
#' @param N Number of vertices / nodes in the network.
#' @param startnodes Initially infectious individuals.
#' @param gam recovery rate
#' @param tau per link infection rate
#' @param graph the adjacency matrix representation of the network.
#' @param status the initial infectious status of individuals.
#' @keywords internal
#' @return Integer array of size \code{N} representing the rates applicable to
#'   individuals
#' @export
init_rates_sis_list <- function(graph, startnodes, gam, tau, status) {
  N <- length(graph)
  rate <- integer(N)
  rate[startnodes] <- gam

  for (i in startnodes) {
    s_contacts <- sus_contacts_list(graph[[i]], status)
    rate[s_contacts] <- rate[s_contacts] + tau
  }

  return(rate)
}


#' Update the simulation status
#'
#' \code{status_update_sis} calculates the new status of a node for SIS update
#'
#' @param status The current infectious status of nodes in graph.
#' @keywords internal
#' @return updated status vector
status_update_sis <- function(status) {
  # change the status of the affected node
  status <- (status + 1) %% 2
  return(status)
}


#' Update the simulation rates
#'
#' \code{rate_update_sis_matrix} calculates the new rate vector for all nodes in
#'   graph
#'
#' @param event The node whose status is changing.
#' @param rate The current rate for each node.
#' @param status The initial infectious status of individuals.
#' @param graph_edge vector of connections to event from adjacency matrix graph
#' @param tau Per link infection rate.
#' @param gam Per node recovery rate.
#' @keywords internal
#' @return updated rate vector
rate_update_sis_matrix <- function(event, rate, status, graph_edge, tau, gam) {
  # find which other susceptible nodes are connected
  s_contacts <- sus_contacts_matrix(graph_edge, status)

  # changes depend on the event
  if (status[event] == 1) {
    # S now an I, so can recover with rate gam
    rate[event] <- gam

    # all contacts now have another I neighbour
    rate[s_contacts] <- rate[s_contacts] + tau
  } else {
    # I now an S
    i_contacts <- inf_contacts_matrix(graph_edge, status)

    # now an S so has rate based on number of infectious neighbours
    rate[event] <- tau * length(i_contacts)

    # all susceptible neighbours lose an I neighbour
    rate[s_contacts] <- rate[s_contacts] - tau
  }
  return(rate)
}


#' Update the simulation rates
#'
#' \code{rate_update_sis_list} calculates the new rate vector for all nodes in
#'   graph
#'
#' @param event The node whose status is changing.
#' @param rate The current rate for each node.
#' @param status The initial infectious status of individuals.
#' @param graph_edge vector of connections to event from adjacency list graph
#' @param tau Per link infection rate.
#' @param gam Per node recovery rate.
#' @keywords internal
#' @return updated rate vector
rate_update_sis_list <- function(event, rate, status, graph_edge, tau, gam) {
  # find which other susceptible nodes are connected
  s_contacts <- sus_contacts_list(graph_edge, status)

  # changes depend on the event
  if (status[event] == 1) {
    # S now an I, so can recover with rate gam
    rate[event] <- gam

    # all contacts now have another I neighbour
    rate[s_contacts] <- rate[s_contacts] + tau
  } else {
    # I now an S
    i_contacts <- inf_contacts_list(graph_edge, status)

    # now an S so has rate based on number of infectious neighbours
    rate[event] <- tau * length(i_contacts)

    # all susceptible neighbours lose an I neighbour
    rate[s_contacts] <- rate[s_contacts] - tau
  }
  return(rate)
}


#' Run a Susceptible -> Infective - Susceptible (SIS) simulation
#'
#' \code{sis_simulation} is a generic function that takes a graph represented
#' either by an adjacency matrix (of class simnetr_matrix) or an adjacency list
#' (of class simnetr_list) and simulates an SIS simulation for a given rate of
#' infection, rate of recovery and duration.
#'
#' @param graph The representation of the network.
#' @param tau Per link infection rate.
#' @param gam Per node recovery rate.
#' @param max_time Maximum simulation runtime.
#' @param i0 Initial number of infective individuals
#' @return list containing timeseries of events and the corresponding timeseries
#'   of number of infectious individual.
#' @examples
#' graph <- erdos_matrix_r(N = 1000, avk = 7)
#' ts <- sis_simulation(graph, tau = 0.5, gam = 1, max_time = 5, i0 = 50, N = 1000)
#' @export
sis_simulation <- function(graph, tau, gam, max_time, i0) {
  UseMethod("sis_simulation")
}



#' Run a Susceptible -> Infective - Susceptible (SIS) simulation
#'
#' \code{sis_simulation.matrix} takes a graph represented by an
#' adjacency matrix and simulates an SIS simulation for
#' a given rate of infection, rate of recovery and duration.
#'
#' @param graph The adjacency matrix representation of the network.
#' @param tau Per link infection rate.
#' @param gam Per node recovery rate.
#' @param max_time Maximum simulation runtime.
#' @param i0 Initial number of infective individuals
#' @return list containing timeseries of events and the corresponding timeseries
#'   of number of infectious individual.
#' @export
sis_simulation.matrix <- function(graph, tau, gam, max_time, i0) {

  # number of nodes in network
  N <- nrow(graph)

  # randomly infect i0 nodes initially
  startnodes <- init_startnodes(N, i0)

  # all nodes initially 'S' (status = 0)
  # except the startnodes which are 'I' (status = 1)
  status <- init_status(N, startnodes)

  # preallocate a vector of rates for each node
  # the initial 'I' nodes have a recovery rate of gam
  # initial 'S' nodes have to increase their rate by tau (S->I) for each
  # of their 'I' neighbours
  rate <- init_rates_sis_matrix(graph, startnodes, gam, tau, status)

  # create empty vectors to store time and number of infectives
  i <- 1
  t <- 0
  inf <- i0

  # start the simulation
  while ((t[i] < max_time) & (inf[i] > 0)) {
    # which node is the vent happening to
    event <- which_event(rate)

    # update values (step and state must be updated prior to rate)
    step <- step_update(rate)
    status[event] <- status_update_sis(status[event])
    rate <- rate_update_sis_matrix(event, rate, status, graph[, event], tau, gam)

    # update time and infectives
    i <- i + 1
    t[i] <- t[i - 1] + step
    inf[i] <- sum(status)
  }

  return(list("t" = t, "inf" = inf))
}



#' Run a Susceptible -> Infective - Susceptible (SIS) simulation
#'
#' \code{sis_simulation.list} takes a graph represented by an
#' adjacency matrix and simulates an SIS simulation for a given rate of
#' infection, rate of recovery and duration.
#'
#' @param graph The adjacency matrix representation of the network.
#' @param tau Per link infection rate.
#' @param gam Per node recovery rate.
#' @param max_time Maximum simulation runtime.
#' @param i0 Initial number of infective individuals
#' @return list containing timeseries of events and the corresponding timeseries
#'   of number of infectious individual.
#' @export
sis_simulation.list <- function(graph, tau, gam, max_time, i0) {
  # number of nodes in network
  N <- length(graph)

  # randomly infect i0 nodes initially
  startnodes <- init_startnodes(N, i0)

  # all nodes initially 'S' (status = 0)
  # except the startnodes which are 'I' (status = 1)
  status <- init_status(N, startnodes)

  # preallocate a vector of rates for each node
  # the initial 'I' nodes have a recovery rate of gam
  # initial 'S' nodes have to increase their rate by tau (S->I) for each
  # of their 'I' neighbours
  rate <- init_rates_sis_list(graph, startnodes, gam, tau, status)

  # create empty vectors to store time and number of infectives
  i <- 1
  t <- 0
  inf <- i0

  # start the simulation
  while ((t[i] < max_time) & (inf[i] > 0)) {
    # which node is the vent happening to
    event <- which_event(rate)

    # update values (step and state must be updated prior to rate)
    step <- step_update(rate)
    status[event] <- status_update_sis(status[event])
    rate <- rate_update_sis_list(event, rate, status, graph[[event]], tau, gam)

    # update time and infectives
    i <- i + 1
    t[i] <- t[i - 1] + step
    inf[i] <- sum(status)
  }

  return(list("t" = t, "inf" = inf))
}