SIR model on the unit square

SIR_randomwalk.R

library(ggplot2)
library(gganimate)
library(dplyr)
rwbounded <- function(U,scaling){
  u <- runif(length(U),-1,1)
  return(U + scaling*ifelse(u>0,(1-U)*u,U*u))
}

rwSIR <- function(N,Tmax,ini,scaling,r,infect_p,remove_p,rwfun=rwbounded,seed=NULL){
  if(is.null(seed)){
    seed <- sample.int(.Machine$integer.max,1)
  }
  set.seed(seed)
  cat("set.seed(",seed,")\n")
  U <- matrix(runif(N*2),N,2)
  out <- data.frame(expand.grid(id=1:N,time=1:Tmax,x=0,y=0,state=0))
  out[out$time==1,3:4] <- U
  for(i in 2:(ini-1)){
    U <- rwfun(U,scaling=scaling)
    out[out$time==i,3:4] <- U
  }
  U <- rwfun(U,scaling=scaling)
  out[out$time==ini,3:4] <- U
  It <- sample.int(N,1)
  s <- integer(N)
  s[It] <- 1L
  out$state[out$id==It & out$time==ini] <- 1L
  for(i in (ini+1):Tmax){
    U <- rwfun(U,scaling=scaling)
    out[out$time==i,3:4] <- U
    w <- s == 1
    if(any(w)){
      adj <- apply(U[w,,drop=FALSE],1,function(pos){rowSums((U[s==0,]-pos)^2) < r})
      adj <- apply(adj, 1, any)
      a <- rbinom(sum(adj),1,infect_p)
      s[s==0][adj] <- a
    }
    b <- rbinom(sum(w),1,remove_p)
    s[w][b==1] <- 2L
    out$state[out$time==i] <- s
  }
  lab <- c("susceptible","infectious","removed")
  out$state <- lab[out$state+1L]
  return(out)
}

out <- rwSIR(N = 100,
      ini = 7,
      Tmax = 100,
      scaling = 0.03,
      r = 0.01,
      infect_p = 0.1,
      remove_p = 0.1,
      seed = 620)

outm <- group_by(out,time,state) %>% 
  tally() %>% 
  mutate(panel_dummy="population")
outf <- rename(outm,time_dummy=time)

outr <- dplyr::filter(out,state!="removed") %>% 
  mutate(panel_dummy="current status")

ggplot(outr,aes(colour=state))+
  geom_point(aes(x=x,y=y),size=2,alpha=0.6)+
  geom_point(data=outm,aes(x=time,y=n),size=2)+
  geom_line(data=outf,aes(x=time_dummy,y=n))+
  facet_wrap(~panel_dummy,scales="free",nrow=2)+
  theme_minimal(14)+theme(legend.position = "bottom")+
  labs(x="",y="")+
  transition_manual(time)

anim_save("./Desktop/SIRanime.gif")