R/run_simulation.R

In qleclerc/epicspatial: Flexible Spatial Modelling of Epidemics

#' @title Runs the spatial epidemic simulation
#'
#' @description Simulates an epidemic using the provided spatial dataset, spatial kernel, contact matrix, and
#'              infection parameters.
#'
#' @param spatial_data The spatial dataset containing the population data.
#' @param D The expanded kernel matrix to use for FOI calculation (generated by the \code{\link{calc_beta}}
#'          function).
#' @param contact_mat The contact matrix for mixing between age groups.
#' @param beta The beta value for the epidemic (calculated from a given R0 using the \code{\link{calc_beta}}
#'             function).
#' @param sigma The recovery rate for the epidemic (must match the one used to calculate beta from R0
#'              using the \code{\link{calc_beta}} function).
#' @param stoch Logical. If TRUE, the simulation is stochastic.
#' @param step Size of time step for stochastic simulation, in days (default is 1 day).
#' @param start_area Where to start the epidemic. Default: area with highest population density. If the spatial
#'                   dataset is a SpatialPolygonsDataFrame, you can specify a string indicating the name of your
#'                   desired starting area. Alternatively, no matter which type of spatial dataset you have, you can
#'                   indicate the index of you desired starting area.
#' @param start_num Number of infected individuals to start the epidemic.
#' @param t_max How many days to run the simulation for.
#'
#' @details This functions requires specific objects to run. These can be generated using the \code{\link{prep_simulation}}
#'          function (e.g. if you want to simulate an epidemic using the spatial dataset "testdata", you must
#'          prep_simulation(test_data) first). The model used is an SIR model, where individuals can be either
#'          Susceptible, Infected or Recovered with regards to the disease. This assumes that Infected individuals
#'          are infectious, and that Recovered individuals are immune and cannot be reinfected.
#'
#' @return Returns one dataframe object containing the estimates of Susceptible, Infected and Recovered individuals
#'         for each time step.
#'
#' @examples
#'
#' #Create a spatial dataset:
#' test_data = raster(nrow=10, ncol=10, xmn=1, xmx=100000, ymn=1, ymx=100000)
#' values(test_data) = runif(100, 1, 1000)
#'
#' #Calculate the parameters for the simulation:
#' prep_simulation(test_data)
#'
#' #Run the simulation:
#' results = run_simulation(test_data, expanded_D, contact_mat, beta)
#'
#' @export


run_simulation = function(spatial_data, D, contact_mat, beta, sigma = 1/2.6, stoch=FALSE, step=1, start_area=NA, start_num=10, t_max=100){

  if(class(spatial_data) == "RasterLayer"){

    populated_areas = which(!is.na(spatial_data@data@values))

    N = spatial_data@data@values[populated_areas]

  } else if(class(spatial_data) == "SpatialPolygonsDataFrame"){

    N = spatial_data@data$population

  } else{

    stop("Incorrect spatial dataset! The dataset must be either a RasterLayer or a SpatialPolygonsDataFrame.")

  }


  num_ages = dim(contact_mat)[1]    #derive number of age categories from contact matrix
  num_areas = dim(D)[1]/num_ages   #derive number of areas from expanded kernel (dim(D)=num_areas*num_ages)



  #### SETTING UP POPULATION: ####

  #N: total population
  #S: susceptibles
  #I: infected
  #R: recovered

  #!!! work in progress, only supports 4 age categories or no age categories at all right now !!!#

  if(num_ages == 4){

    N = matrix(N, nrow=num_areas, ncol=num_ages)

    #set minimum population in an area to 1:
    N[which(N<1)] = 1

    NN0 = as.vector(N)

    N[,1] = N[,1]*(5/81)
    N[,2] = N[,2]*(14/81)
    N[,3] = N[,3]*(46/81)
    N[,4] = N[,4]*(16/81)
    #this way, i in the matrix is the area and j the age group e.g. N[1,1] gives pop 0-4 in area 1

    S = N
    I = matrix(0, nrow=num_areas, ncol=num_ages)
    R = I


  } else if(num_ages == 1){

    N[which(N<1)] = 1

    NN0 = as.vector(N)

    S = N
    I = rep(0, num_areas)
    R = I

  } else {

    stop("Unsupported number of age categories, currently only supports 4 or none.")

  }



  #### SEEDING: ####

  #seeds by making a given number of adults infected in the chosen starting area

  if(is.na(start_area)){

    #identify area with the most inhabitants:

    if(num_ages == 1){

      start_area = which.max(N)

    } else {

      #which area has the highest total population? (rowSums gives total pop in each area since each column is an age group)
      start_area = which.max(rowSums(N))

    }

  }

  else if(is.numeric(start_area)){

    #extract start area based on index number:
    start_area = start_area

  }

  else{

    #extract area based on name:

    start_area = which(spatial_data@data$name == start_area)

  }



  #change pop values in starting area:
  if(num_ages == 1){

    S[start_area] = S[start_area] - start_num
    I[start_area] = I[start_area] + start_num

  } else {

    S[start_area,3] = S[start_area,3] - start_num
    I[start_area,3] = I[start_area,3] + start_num
    #3 is the adult age group

  }



  #### RUN MODEL: ####

  #deterministic model:
  full_model = function(t, y, .) {

    #counter to display progress during model execution:
    print(paste0(floor(t/t_max*100), "% done"))

    S = y[1:(num_areas*num_ages)]
    I = y[(num_areas*num_ages+1):(num_areas*num_ages*2)]
    R = y[(num_areas*num_ages*2+1):(num_areas*num_ages*3)]

    lambda=beta*S*(D %*% (I/NN0))

    dSdt = -lambda
    dIdt = lambda - sigma*I
    dRdt = sigma*I

    list(c(dSdt, dIdt, dRdt))

  }


  #stochastic model:
  if(stoch == TRUE){

    #inverse step for cleaner calculations involving this value below:
    step = 1/step

    S = round(S)
    S[which(S==0)] = 1

    S = as.vector(S)
    I = as.vector(I)
    R = as.vector(R)

    #adjust beta and sigma values according to time step:
    beta = beta/step
    sigma = sigma/step

    #create matrix to store the results:
    results = matrix(0, (t_max*step+1), (num_areas*num_ages*3+1))

    #adjust the values of the "Time" column:
    results[,1] = seq(0,t_max*step,1)

    #add the starting population values:
    results[1,-1] = c(S,I,R)

    #convert sigma from a rate to a probability:
    sigma = 1 - exp(-(sigma))

    for(t in 1:(t_max*step)){

      lambda = beta*(D %*% (I/NN0))
      lambda = 1 - exp(-lambda)

      for(i in 1:length(N)){

        new_inf = rbinom(1,S[i],lambda[i])
        rec = rbinom(1, I[i], sigma)

        S[i] = S[i] - new_inf
        I[i] = I[i] + new_inf - rec
        R[i] = R[i] + rec

      }

      #store results:
      results[t+1,-1] = c(S,I,R)

    }


  } else {

    #set method to rk4 rather than default lsoda
    #required because otherwise deSolve requires too much RAM
    #more consistent performance, but slower when smaller number of areas
    #parms=0 because rk4 requires a definition of parms, but effectively the function inherits the parameters from the parent environment

    results = deSolve::ode(method="rk4", func=full_model, y=c(S,I,R), times=seq(0,t_max,1), parms=0)

  }


  return(results)

}