SIR model on the lattice

SIRgrid.R

library(dplyr)
library(tidyr)
library(gganimate)

SIRsim <- function(beta,gamma,nx,ny,nt,seed=NULL){
  if(is.null(seed)){
    set.seed(seed)
  }
  mat <- matrix(0L,nx+2,ny+2)
  mat[sample.int(nx,1)+1L,sample.int(ny,1)+1L] <- 1L
  out <- array(0L,c(nx,ny,nt))
  out[,,1] <- mat[2:(nx+1),2:(ny+1)]
  for(k in 2:nt){
    for(i in 2:(nx+1)){
      for(j in 2:(ny+1)){
        #I ->R
        if(mat[i,j]==1){mat[i,j] <- rbinom(1,1,gamma)+1}
        #S -> I
        if(mat[i-1,j]==1 & mat[i,j]==0){mat[i,j] <- rbinom(1,1,beta)}
        if(mat[i+1,j]==1 & mat[i,j]==0){mat[i,j] <- rbinom(1,1,beta)}
        if(mat[i,j-1]==1 & mat[i,j]==0){mat[i,j] <- rbinom(1,1,beta)}
        if(mat[i,j+1]==1 & mat[i,j]==0){mat[i,j] <- rbinom(1,1,beta)}
      }
    }
    out[,,k] <- mat[2:(nx+1),2:(ny+1)]
  }
  return(out)
}

res <- SIRsim(0.05,0.01,20,20,400,123)
resdf <- expand.grid(x=1:dim(res)[1],y=1:dim(res)[2], time=1:dim(res)[3]) %>% 
  as.data.frame() %>% 
  mutate(state = as.vector(res)) %>% 
  mutate(state = factor(c("S","I","R")[state+1L], levels=c("S","I","R"))) %>% 
  mutate(panel_dummy="lattice")

resdf_m <- group_by(resdf,time,state) %>% 
  tally() %>% ungroup() %>% 
  mutate(panel_dummy="marginal")

resdf_m2 = rename(resdf_m,time2=time,n2=n)

ggplot(resdf)+
  geom_point(data=resdf_m,aes(x=time,y=n,colour=state),size=2)+
  geom_line(data=resdf_m2, aes(x=time2,y=n2,colour=state))+
  geom_tile(aes(x=x,y=y,fill=state),colour="black")+
  facet_wrap(~panel_dummy,scales="free")+
  scale_fill_manual(values = c("cornflowerblue","salmon","gray"))+
  scale_colour_manual(values = c("cornflowerblue","salmon","gray"))+
  labs(x="",y="")+
  theme_minimal(14)+theme(legend.position = "bottom")+
  transition_manual(time)

anim_save("~/Desktop/SIRlattice.gif")