SIR model on the unit square

SIR_randomwalk.R

library(ggplot2)
library(gganimate)
library(dplyr)

rwSIR <- function(N,Tmax,ini,scaling,r2,infect_p,remove_p,
                  randu = function(U){runif(length(U),-1,1)},
                  decay=1,seed=NULL){
  rwbounded1 <- function(U,scaling){
    u <- randu(U)
    scaling*ifelse(u>0,(1-U)*u,U*u)
  }
  rwbounded2 <- function(U,scaling,u0,decay){
    u <- randu(U)
    u <- (1-decay)*u0 + decay*u
    scaling*ifelse(u>0,(1-U)*u,U*u)
  }
  binarysample <- function(n,p){
    sample(c(TRUE,FALSE),size=n,replace = TRUE,prob = c(p,1-p))
  }
  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
  u0 <- rwbounded1(U,scaling=scaling)
  for(i in 2:(ini-1)){
    u0 <- rwbounded2(U,scaling,u0,decay)
    U <- U + u0
    out[out$time==i,3:4] <- U
  }
  u0 <- rwbounded2(U,scaling,u0,decay)
  U <- U + u0
  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){
    u0 <- rwbounded2(U,scaling,u0,decay)
    U <- U + u0
    out[out$time==i,3:4] <- U
    s1 <- s == 1
    s0 <- s == 0
    if(any(s1)&any(s0)){
      ad <- apply(U[s1,,drop=FALSE],1,function(pos){
        a <- rowSums((U[s0,,drop=FALSE]-pos)^2) < r2
        a[a] <- binarysample(sum(a),infect_p)
        return(a)
      })
      if(!is.null(dim(ad))){
        ad <- apply(ad,1,any)
      }
      s[s0][ad] <- 1L
    }
    b <- binarysample(sum(s1),remove_p)
    s[s1][b] <- 2L
    out$state[out$time==i] <- s
  }
  lab <- c("susceptible","infectious","removed")
  out$state <- lab[out$state+1L]
  return(out)
}
scl <- runif(500,0,0.1)
out <- rwSIR(N = 500,
      ini = 7,
      Tmax = 200,
      scaling = scl,
      r2 = 0.001,
      infect_p = 0.2,
      remove_p = 0.1,
      decay = 0.5,
      seed = 620)

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

outf <- rename(outm,time_dummy=time)

outr <- dplyr::filter(out) %>% 
  mutate(panel_dummy=if_else(state=="removed","B","A"))

ggplot(outr,aes(colour=state))+
  geom_point(aes(x=x,y=y),size=2,alpha=0.8)+
  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")+
  theme_minimal(14)+theme(legend.position = "bottom")+
  labs(x="",y="")+
  transition_manual(time)

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