contagion modeling with mobility O/D matrices

covid-mobility-sir.R

library(tidyverse)
library(data.table)
library(glue)
library(parallel)
library(pbmcapply)
here <- here::here

OD <- read_csv(here("output/OD.csv"))
OD.norm <- t(t(OD) / colSums(OD))
pop <- read_csv(here("output/pop.csv"), col_types="nc")

FIRST_INFECT_THRESHOLD = 600
N_CORES = 16

SIR <- data.frame(
  S=pop$pop,
  I=rep(0, nrow(pop)),
  R=rep(0, nrow(pop)))
rownames(SIR) <- pop$sa1

# 1% of each SA1 over FIRST_INFECT_THRESHOLD starts out infected
infections.first <- (SIR$S >= FIRST_INFECT_THRESHOLD) * round(SIR$S / 100)

# update SIR columns

SIR$S = SIR$S - infections.first
SIR$I = SIR$I + infections.first

SIR.norm <- SIR / rowSums(SIR)

## Parameters from
## https://docs.google.com/spreadsheets/d/1YEj4Vr6lG1jQ1R3LG6frijJYNynKcgTjzo2n0FsBwZA/htmlview?#
## and
## https://www.nature.com/articles/s41421-020-0148-0#Sec2

# average recovery as 18 days
gamma <- 1/18
R0 <- 3.1
# according to some reports, R0 is conservatively estimated to be 2.5.
# this means that 2.5*0.04 = beta


#mean_beta <- 1/2.5
mean_beta <- R0*gamma
betas <- rgamma(nrow(pop), shape=mean_beta/2, scale=2)
gammas <- rep(gamma, nrow(pop))

#R0 = mean(betas)/gamma

# Modeling
# ===================================

  
run_model <- function (n_steps, flow_pct=1, save=F, outpath="output/SIR",print=F) {
  times <- seq(0, n_steps)
  SIR.sim <- copy(SIR)
  SIR.norm.sim <- copy(SIR.norm)

  OD.exp <- OD
  OD.exp.norm <- t(t(OD.exp) / colSums(OD.exp))

  runs <- lapply(times, function (t) {
    ## total # of flow that is infected.
    infected.OD.exp <- SIR.norm.sim$I * as.matrix(OD.exp)
    ## total inflow that is infected
    ## FLOW_PCT here.
    infected.inflow <- round(flow_pct * colSums(SIR.norm.sim$I * as.matrix(OD.exp)))
    
    ## print(glue("{t} : \t total inflow: {sum(infected.inflow)}"))
    
    ## introductions to places w/ no current infections is proportional to inflow.
    infected.inflow.norm <- colSums(SIR.norm.sim$I * as.matrix(OD.exp.norm))
    hazard <- (betas * SIR.norm.sim$S) * (1 - (exp(-infected.inflow.norm) * SIR.norm.sim$S)) / (1 + betas * SIR.norm.sim$S)
    infected.introductions <- sapply(hazard, function (h) rbinom(1, 1, h)) * (SIR.sim$I == 0)
    SIR.sim$S <<- SIR.sim$S - infected.introductions
    SIR.sim$I <<- SIR.sim$I + infected.introductions
    
    # where disease has been introduced, use in-district trips and inflow 
    # from other districts
    #infected.new <- round(betas * SIR.sim$S * SIR.sim$I / pop$pop) 
    infected.inflow <- round((betas * SIR.sim$S * infected.inflow) / (pop$pop + colSums(OD.exp)))
    
    ## cap it to the total number of susceptible.
    infected.total <- infected.inflow #+ infected.new)
    infected.total <- (infected.total <= SIR.sim$S) * infected.total
    infected.total <- infected.total + ((infected.total > SIR.sim$S) * SIR.sim$S)

    ## recovered population    
    recovered.new <- round(SIR.sim$I * gammas)
    
    S = sum(SIR.sim$S) / sum(pop$pop)
    I = sum(SIR.sim$I) / sum(pop$pop)
    R = sum(SIR.sim$R) / sum(pop$pop)

    if (print) {
      print(glue("{t} : \t{S}\t{I}\t{R}\t{sum(pop$pop)}"))      
    }    
    
    SIR.sim$S <<- SIR.sim$S - infected.total
    SIR.sim$I <<- SIR.sim$I + infected.total - recovered.new
    SIR.sim$R <<- SIR.sim$R + recovered.new
    
    SIR.norm.sim <<- SIR.sim / rowSums(SIR.sim)
    
    data.frame(sa1=pop$sa1, 
               S=SIR.sim$S, 
               I=SIR.sim$I, 
               R=SIR.sim$R,
               t=t)    
  })  
  
  runs.all <- rbindlist(runs)
  
  runs.norm <- runs.all %>%
    group_by(t) %>%
    select(-sa1) %>%
    summarise_all(sum)
  runs.norm[,2:4] <- runs.norm[,2:4] / rowSums(runs.norm[,2:4])
  
  runs.all$flow.pct <- flow_pct
  runs.norm$flow.pct <- flow_pct
  
  if (save) {
    write_csv(runs.norm, 
              glue("{outpath}/SIR_alpha-{flow_pct}_beta-{beta}_gamma-{gamma}.csv"))
    write_csv(runs.all, 
              glue("{outpath}/SIR_raw_alpha-{flow_pct}_beta-{beta}_gamma-{gamma}.csv"))
  }
  
  list(raw=runs.all,
       norm=runs.norm)
}



## run our model, with alpha varying from 1 to 0.1.
## -----------------------------------------------------------
flow_pcts <- seq(1, 0.1, -0.1)

dir.create("output/SIR", showWarnings=F)
runs.flows <- pbmclapply(flow_pcts, run_model, n_steps=150, save=T, mc.cores=N_CORES)


## Plotting
## ====================================

plot.SIR <- function (runs.norm) {
  runs.norm %>%
    pivot_longer(cols=c(S, I, R)) %>%
    ggplot(aes(x=t, color=name, y=value)) +
    geom_path()
}

#plots <- lapply(runs.flows, function (f) plot.SIR(f$norm))