SIR model(1d)

SIR_randomwalk.R

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

rwSIR1d <- function(N,Tmax,ini,scaling,r,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 <- runif(N)
  out <- data.frame(expand.grid(id=1:N,time=1:Tmax,x=0,state=0))
  out$p[out$time==1] <- U
  u0 <- rwbounded1(U,scaling=scaling)
  for(i in 2:(ini-1)){
    u0 <- rwbounded2(U,scaling,u0,decay)
    U <- U + u0
    out$x[out$time==i] <- U
  }
  u0 <- rwbounded2(U,scaling,u0,decay)
  U <- U + u0
  out$x[out$time==ini] <- 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$x[out$time==i] <- U
    s1 <- s == 1
    s0 <- s == 0
    if(any(s1)&any(s0)){
      ad <- sapply(U[s1],function(pos){
        a <- abs(U[s0]-pos) < r
        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(100,0,0.1)
out <- rwSIR1d(N = 100,
      ini = 7,
      Tmax = 200,
      scaling = scl,
      r = 0.01,
      infect_p = 0.2,
      remove_p = 0.1,
      decay = 0.5,
      seed = 620)

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

ggplot(outm,aes(x=time,y=n,colour=state))+
  geom_line()+
  theme_minimal(14)+theme(legend.position = "bottom")+
  labs(x="",y="")

outf <- rename(outm,time_dummy=time)

#out <- mutate(out,panel_dummy="current status")
outr <- dplyr::filter(out) %>% 
  mutate(panel_dummy="removed",y_dummy=as.integer(state=="removed"))

ggplot(outr,aes(colour=state))+
  geom_point(aes(x=x,y=y_dummy),size=5,alpha=0.3)+
  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")