A simulation of viral transmission, and various measures to slow the spread

flatten_the_curve.R

library(tidyverse)
library(colourlovers)
treatment_colors <- swatch(clpalette("1473"))[[1]]

sim_function <- function(
  n_population = 10000,
  n_days = 100,
  time_to_quarantine = 14,
  n_encounters = 50,
  transmission_probability = 0.05,
  n_initial = 1
){
  condition_matrix<- matrix(nrow = n_population,
                            ncol = n_days,
                            0)
  condition_matrix[1:n_initial, 1] <- 1
  condition_counts <- matrix(nrow = 3, ncol = n_days, NA)
  current_count <- c(table(condition_matrix[, 1])[c("-1", "0", "1")])
  current_count <- replace_na(current_count, 0)
  condition_counts[, 1] <- current_count
  colnames(condition_counts) <- 1:n_days
  rownames(condition_counts) <- c(-1, 0, 1)
  
  for(jj in 2:ncol(condition_matrix)){
    prop_infected <- mean(condition_matrix[, jj-1] == 1)
    
    # Number of encounters with infected
    expected_encounters_with_infected <- n_encounters * prop_infected
    n_encounters_with_infected <- rbinom(n_population, n_encounters, prop_infected)
    
    # Whether individual contracted virus
    prop_did_not_contract <- (1 - transmission_probability) ^ n_encounters_with_infected
    contracted <- rbinom(n_population, 1, 1-prop_did_not_contract)
    condition_matrix[, jj] <- contracted
    
    # Plus anyone who already had it
    condition_matrix[, jj][condition_matrix[, jj-1] == 1] <- 1
    condition_matrix[, jj][condition_matrix[, jj-1] == -1] <- -1
    
    # Go under quarantine
    condition_recognized <- rowSums(condition_matrix[, 1:jj] == 1) >= time_to_quarantine
    condition_matrix[condition_recognized, jj] <- -1
    
    current_count <- c(table(condition_matrix[, jj])[c("-1", "0", "1")])
    current_count <- replace_na(current_count, 0)
    condition_counts[, jj] <- current_count
  }
  
  n_infected_by_day <- colSums(condition_counts[-2, ])
  n_newly_infected <- c(0, colSums(condition_matrix[, -1] == 1 & condition_matrix[, -ncol(condition_matrix)] == 0))
  
  tibble(day = 1:n_days,
         n_infected_by_day,
         n_newly_infected) %>%
    return()
}

test <- sim_function()

# Parameter sweep
scenarios <- expand.grid(iteration = 1:25,
                         time_to_quarantine = c(7, 14),
                         n_encounters = c(50, 100),
                         transmission_probability = c(0.005, 0.01))
scenarios <- scenarios %>%
  mutate(quarantine = as.factor(time_to_quarantine),
         distancing = as.factor(n_encounters),
         washing = as.factor(transmission_probability))
levels(scenarios$quarantine) <- c("testing", "symptoms")
levels(scenarios$distancing) <- c("distancing", "not distancing")
levels(scenarios$washing) <- c("washing", "not washing")

# Simulate
results <- scenarios %>%
  group_by(iteration,
           quarantine, distancing, washing) %>%
  do(with(.,
          sim_function(
            n_population = 10000,
            time_to_quarantine = time_to_quarantine,
            n_encounters = n_encounters,
            transmission_probability = transmission_probability))) %>%
  ungroup()

results <- results %>% sample_frac(1)

mean_by_day <- results %>%
  group_by(quarantine, distancing, washing, day) %>%
  dplyr::summarise(n_infected_by_day = median(n_infected_by_day),
                   n_newly_infected = median(n_newly_infected)) %>%
  ungroup()

ggplot(mean_by_day) +
  aes(x = day, y = n_infected_by_day,
      colour = paste(quarantine, distancing, washing)) +
  geom_line(lwd = 2) +
  #scale_colour_discrete(guide = "none") +
  scale_colour_brewer(type = "qual") +
  scale_y_continuous(expand = c(0, 0)) +
  theme_bw()


ggplot(results) +
  aes(x = day, y = n_infected_by_day,
      colour = paste(quarantine, distancing, washing),
      group = paste(iteration, quarantine, distancing, washing)) +
  geom_line() +
  scale_colour_discrete(guide = "none") +
  theme_bw()

line_alpha <- 5/7
line_lwd <- 1/3
plot_title <- "Flattening the curve"
plot_subtitle <- "Simulating the incidence of viral infection"
plot_caption <- "Simulations with and without implementation of three precautionary measures"

ggplot(results) +
  aes(x = day, y = n_newly_infected,
      colour = distancing,
      group = paste(iteration, quarantine, distancing, washing)) +
  geom_line(alpha = line_alpha, lwd = line_lwd) +
  facet_grid(washing ~ quarantine) +
  scale_colour_manual(values = c(treatment_colors[5], treatment_colors[2])) +
  scale_y_continuous("New infections",
                     expand = expand_scale(mult = c(0, 0.05))) +
  scale_x_continuous("Day", expand = c(0, 0)) +
  labs(title = plot_title,
       subtitle = plot_subtitle,
       caption = plot_caption) +
  guides(col = guide_legend(override.aes = list(shape = 15, size = 6, alpha = 1))) +
  theme_bw()
ggsave(plot = last_plot(), "ftc_distancing.png", h = 6.18, w = 10, type = "cairo-png")

ggplot(results) +
  aes(x = day, y = n_newly_infected,
      colour = washing,
      group = paste(iteration, quarantine, distancing, washing)) +
  geom_line(alpha = line_alpha, lwd = line_lwd) +
  facet_grid(distancing ~ quarantine) +
  scale_colour_manual(values = c(treatment_colors[1], treatment_colors[2])) +
  scale_y_continuous("New infections",
                     expand = expand_scale(mult = c(0, 0.05))) +
  scale_x_continuous("Day", expand = c(0, 0)) +
  labs(title = plot_title,
       subtitle = plot_subtitle,
       caption = plot_caption) +
  guides(col = guide_legend(override.aes = list(shape = 15, size = 6, alpha = 1))) +
  theme_bw()
ggsave(plot = last_plot(), "ftc_washing.png", h = 6.18, w = 10, type = "cairo-png")

ggplot(results) +
  aes(x = day, y = n_newly_infected,
      colour = quarantine,
      group = paste(iteration, quarantine, distancing, washing)) +
  geom_line(alpha = line_alpha, lwd = line_lwd) +
  facet_grid(washing ~ distancing) +
  scale_colour_manual(values = c(treatment_colors[3], treatment_colors[2])) +
  scale_y_continuous("New infections",
                     expand = expand_scale(mult = c(0, 0.05))) +
  scale_x_continuous("Day", expand = c(0, 0)) +
  labs(title = plot_title,
       subtitle = plot_subtitle,
       caption = plot_caption) +
  guides(col = guide_legend(override.aes = list(shape = 15, size = 6, alpha = 1))) +
  theme_bw()
ggsave(plot = last_plot(), "ftc_quarantine.png", h = 6.18, w = 10, type = "cairo-png")

# Logistic plots
mean_by_day <- results %>%
  group_by(quarantine, distancing, washing, day) %>%
  dplyr::summarise(n_infected_by_day = median(n_infected_by_day),
                   n_newly_infected = median(n_newly_infected)) %>%
  ungroup()
levels(mean_by_day$quarantine) <- c("test", "~test")
levels(mean_by_day$distancing) <- c("distance", "~distance")
levels(mean_by_day$washing) <- c("wash", "~wash")

ggplot(mean_by_day) +
  aes(x = day, y = n_infected_by_day,
      colour = paste(quarantine, distancing, washing)) +
  geom_line(lwd = 1, lty = 1, alpha = 1) +
  scale_colour_brewer("Measures taken",
                      #type = "qual",
                      palette = "Dark2") +
  scale_y_continuous("Cumulative infections",
                     expand = expand_scale(mult = c(0, 0.05))) +
  scale_x_continuous("Day", expand = c(0, 0)) +
  labs(title = plot_title,
       subtitle = plot_subtitle,
       caption = plot_caption) +
  guides(col = guide_legend(override.aes = list(shape = 15, size = 6, alpha = 1))) +
  theme_bw()
ggsave(plot = last_plot(), "ftc_cumulative.png", h = 6.18, w = 10, type = "cairo-png")