Simulating disease spread in a population

Question

I am trying to build a model that allows me to estimate the spread of disease in the population. The setup is similar to the SIR model, with the difference that this model includes the stochastic component.

I have written an algorithm to carry out the simulation, but the lead times are really high. Could you please help me optimize the code? Below is the code.

Population = 100000
Time = 365
r_0 = 0.02
start_positive = 100
min_positive_days = 5
max_positive_days = 15

Pop_over_time <- matrix("S", ncol = Time, nrow = Population, dimnames = list(1:Population, paste0("T",1:Time)))
I0 <- sample(Population, start_positive)

for (i in I0) {
  Pop_over_time[i, 1] <- "I"
  nr <- trunc(runif(1, min_positive_days, max_positive_days))
  Pop_over_time[i, 1:nr] <- "I"
  Pop_over_time[i, (nr):(Time)] <- "R"
}


for (i in 2:Time) {
  r_s <- r_0 * sqrt(2 + sin(500 + i * 2 * pi / 365)) / 2
  
     
  for (j in 1:Population) {
    if (Pop_over_time[j, i - 1] == "I") {
      for (z in (1:k[j])) {
        if (Pop_over_time[Link[j, z], i - 1] == "S") {
          Pop_over_time[Link[j, z], i] <- rbinom(1, 1, r_s)
          
          if (Pop_over_time[Link[j, z], i] == 1) {
            nr <- trunc(runif(1, min_positive_days, max_positive_days))
            nr_min <- min((i + nr - 1), Time)
            Pop_over_time[Link[j, z], i:nr_min] <-
              rep(1, length(i:nr_min))
            Pop_over_time[Link[j, z], (nr_min):(Time)] <- "R"
            
          }
        }
      }
    }
    
  }
  Pop_over_time[, i][Pop_over_time[, i] == 1] <- "I"
  Pop_over_time[, i][Pop_over_time[, i] == 0] <- "S"
  
}

In a first part I create a matrix (Link) that defines all the links within the population. So first I define variables, then I create a matrix that saves, for each day, the situation of the population. Then I define the first 100 cases (with a sampling) and after that the heaviest part of the code starts. r_s (the probability of contagion) is defined in a seasonal way. So, for each member of the population that I infect I generate a binomial variable for each of his contacts: if the number generated is 1 then the infection has occurred, vice versa not.

Here you can download the .rdata file to run the code.

Answer 1

1. we should remove dimnames from matrix, as it slows sub-setting. They can be added after loops with: dimnames(Pop_over_time) <- list(1:Population, paste0("T",1:Time))
2. working with characters is slower than integers, so we can replace S/I/R with integers that doesn't interrupt with your current code
3. rbinom isn't as fast when called millions of times inside for loops. So we can move that one level up, and then in loop get relevant value. Maybe this could be improved even more...
4. other minor changes

Code:

S = 4L
I = 5L
R = 6L
Pop_over_time <- matrix(S, ncol = Time, nrow = Population
                        # , dimnames = list(1:Population, paste0("T",1:Time))
)

for (i in I0) {
  Pop_over_time[i, 1] <- I
  nr <- trunc(runif(1, min_positive_days, max_positive_days))
  Pop_over_time[i, 1:nr] <- I
  Pop_over_time[i, (nr):(Time)] <- R
}

for (i in 2:Time) {
  r_s <- r_0 * sqrt(2 + sin(500 + i * 2 * pi / 365)) / 2
  v1 <- Pop_over_time[, i - 1]
  for (j in 1:Population) {
    if (v1[j] == I) {
      rbin2 <- rbinom(k[j], 1, r_s)
      for (z in (1:k[j])) {
        x1 <- Link[j, z]
        if (v1[x1] == S) {
          Pop_over_time[x1, i] <- rbin2[z]
          if (Pop_over_time[x1, i] == 1) {
            nr <- trunc(runif(1, min_positive_days, max_positive_days))
            nr_min <- min((i + nr - 1), Time)
            Pop_over_time[x1, i:nr_min] <- rep(1, length(i:nr_min))
            Pop_over_time[x1, (nr_min):(Time)] <- R
          }
        }
      }
    }
  }
  Pop_over_time[, i][Pop_over_time[, i] == 1] <- I
  Pop_over_time[, i][Pop_over_time[, i] == 0] <- S
}