Cacher is the code snippet organizer for pro developers

rm(list = ls())
p_load(deSolve)

# Thanh Le Viet
# May, 2016
# Human Metapopulation
# http://modelinginfectiousdiseases.org/
# Ported from python version at
# http://homepages.warwick.ac.uk/~masfz/ModelingInfectiousDiseases/Chapter7/Program_7.2/Program_7_2.py

n <- 5 #Number of population

# Setup the state and parameters
N0 <- S0 <- rep(0, n*n)
steps <- seq(1,n*n,n+1)
N0[steps] <- 1000
S0[steps] <- 800
I0 <- c(1,rep(0,n*n-1))

ll <- matrix(rep(0,n*n),nrow = n)
for (i in seq_len(n))
  for (j in seq_len(n))
    if (abs(i-j)==1)
      ll[i,j] <- 0.1

beta <- rep(1,n)*1 #Homogeneous

# beta <- rpois(n,2) #Inhomogeneous

gamma <- rep(1,n)*0.3
r <- matrix(rep(2,n*n),nrow = n)
diag(r) <- 0

INPUT0 <- c(S0,I0,N0)

INPUT <- rep(0, 3*n*n)

INPUT[seq(1,n*n*3,3)]  <- INPUT0[seq(1,n*n,1)]
INPUT[seq(2,n*n*3,3)] <- INPUT0[seq(n*n+1,2*n*n,1)]
INPUT[seq(3,n*n*3,3)] <- INPUT0[seq(n*n*2+1,3*n*n,1)]

metapop <- function(t, state, parameters){
  with(parameters, {
    Y <- rep(0,3*N*N)
    V <- state
    sumY <- rep(0,N)
    sumN <- rep(0,N)

    # R style: try to reduce the for loop and vectorise
    # for (i in seq_len(N)){
    #   k <- seq(1,3*N*N,3*N)
    #   iI <- seq(k[i]+1,3*N+k[i],3)
    #   iN <- seq(k[i]+2,3*N+k[i],3)
    #   sumY[i] <- sum(V[iI])
    #   sumN[i] <- sum(V[iN])
    # }

    # The following block returns the results as same as the block above
    for (I in seq_len(N)){
      for (J in seq_len(N)){
        i <- I - 1
        j <- J - 1
        k <- 3 * (j + i*N) + 1
        sumY[I] <- sumY[I] + V[1+k]
        sumN[I] <- sumN[I] + V[2+k]
      }
    }


    for (I in seq_len(N)){
      for (J in seq_len(N)){
        i <- I - 1
        j <- J - 1
        #Rate
        k <- 3*(j+i*N) + 1
        K <- 3*(i+j*N) + 1
        h <- 3*(i+i*N) + 1
        H <- 3*(j+j*N) + 1

        Y[k] <-  Y[k] - beta[I]*V[k]*(sumY[I]/sumN[I])
        Y[k+1] <-  Y[k+1] + beta[I]*V[k]*(sumY[I]/sumN[I])
        Y[k+1] <- Y[k+1] - gamma[I]*V[k+1]

        #Movement
        Y[h] <-  Y[h] + r[J,I]*V[K]
        Y[h] <- Y[h] - ll[J,I]*V[h]

        Y[h+1] <-  Y[h+1] + r[J,I]*V[K+1]
        Y[h+1] <- Y[h+1] - ll[J,I]*V[h+1]

        Y[h+2] <- Y[h+2] + r[J,I]*V[K+2]
        Y[h+2] <- Y[h+2] - ll[J,I]*V[h+2]

        Y[k] <- Y[k] + ll[I,J]*V[H]
        Y[k] <- Y[k] - r[I,J]*V[k]

        Y[1+k] <- Y[1+k] + ll[I,J]*V[1+H]
        Y[1+k] <- Y[1+k] - r[I,J]*V[1+k]

        Y[2+k] <- Y[2+k] + ll[I,J]*V[2+H]
        Y[2+k] <- Y[2+k] - r[I,J]*V[2+k]
      }
    }
    list(Y)
  })
}
#Set up
parameters <- list(N = n, ll = ll, r = r)
state <- INPUT
ND <- 80
t <- seq(0,ND,1)
# Perform the model
meta_outcome <- as.data.frame(ode(y = state, times = t, func = metapop, parms = parameters))
times <- meta_outcome[,1]
meta_outcome <- meta_outcome[,-1]
# Plot function
plot.meta <- function(meta_outcome, times){
  totalS <- matrix(0, nrow = length(times), ncol = n)
  totalI <- matrix(0, nrow = length(times), ncol = n)

  for (I in seq_len(n)){
    for (J in seq_len(n)){
      i <- I - 1
      j <- J - 1
      k <- 3*(j + i*n) + 1
      totalS[,I] <- totalS[,I] + meta_outcome[,k]
      totalI[,I] <- totalI[,I] + meta_outcome[,k+1]
    }
  }
  I.y.lim <- max(totalI)
  S.y.lim <- max(totalS)
  par(mfrow = c(2,1))
  plot(1, type="n", xlab="", ylab="", xlim=c(0, max(times)), ylim=c(0, S.y.lim), main = paste0("S, pop = ",n))
  for (i in seq(n)){
    lines(times, totalS[,i], col = i)
  }
  plot(1, type="n", xlab="", ylab="", xlim=c(0, max(times)), ylim=c(0, I.y.lim), main = "I")
  for (i in seq(n)){
    lines(times, totalI[,i], col = i)
  }
}
#Plot outcome
plot.meta(meta_outcome = meta_outcome, times = times)