Individual-based modeling in Julia

SimJulia is a package for process-based simulations, similar to SimPy in Python, which I contributed to a little to many moons ago. Over the years, I've implemented simulation models of disease transmission to help me learn new computer languages (see my efforts in Python, C, and Eiffel on my Google Code site here). Here's my first try at a process-based version of a stochastic SIR type epidemiological model.

module SIRSim

using SimJulia
using Distributions
using DataFrames

export
    SIRPerson,
    SIRModel,
    activate,
    run,
    out

type SIRPerson
    state::Char
end

function increment(a::Array{Int64})
    push!(a,a[length(a)]+1)
end

function decrement(a::Array{Int64})
    push!(a,a[length(a)]-1)
end

function carryover(a::Array{Int64})
    push!(a,a[length(a)])
end

type SIRModel
    sim::Simulation
    parray::Array{Process}
    beta::Float64
    c::Float64
    gamma::Float64
    ta::Array{Float64}
    Sa::Array{Int64}
    Ia::Array{Int64}
    Ra::Array{Int64}
    allIndividuals::Array{SIRPerson}
end

function SIRModel(heapsize::Int64,beta::Float64,c::Float64,gamma::Float64,S::Int64,I::Int64,R::Int64)
        sim = Simulation(uint(heapsize))
        N=S+I+R
        parray = [Process(sim,string(i)) for i in 1:N]
        states = [fill!(Array(Char,S),'S'),
                  fill!(Array(Char,I),'I'),
                  fill!(Array(Char,R),'R')]
        allIndividuals=[SIRPerson(state) for state in states]
        ta=Array(Float64,0)
        push!(ta,0.0)
        Sa=Array(Int64,0)
        push!(Sa,S)
        Ia=Array(Int64,0)
        push!(Ia,I)
        Ra=Array(Int64,0)
        push!(Ra,R)
        SIRModel(sim,parray,beta,c,gamma,ta,Sa,Ia,Ra,allIndividuals)
end


function live(p::Process,individual::SIRPerson,s::SIRModel)
  while individual.state=='S'
      # Wait until next contact
      hold(p,rand(Distributions.Exponential(1/s.c)))
      # Choose random alter
      alter=individual
      while alter==individual
          N=length(s.allIndividuals)
          index=rand(Distributions.DiscreteUniform(1,N))
          alter=s.allIndividuals[index]
      end
      # If alter is infected
      if alter.state=='I'
          infect = rand(Distributions.Uniform(0,1))
          if infect < s.beta
              individual.state='I'
              push!(s.ta,now(p))
              decrement(s.Sa)
              increment(s.Ia)
              carryover(s.Ra)
          end
      end
  end
  if individual.state=='I'
      # Wait until recovery
      hold(p,rand(Distributions.Exponential(1/s.gamma)))
      individual.state='R'
      push!(s.ta,now(p))
      carryover(s.Sa)
      decrement(s.Ia)
      increment(s.Ra)
  end
end

function activate(s::SIRModel)
    [SimJulia.activate(s.parray[i],0.0,live,s.allIndividuals[i],s) for i in 1:length(s.parray)]
end

function run(s::SIRModel,tf::Float64)
    SimJulia.run(s.sim,tf)
end

function out(s::SIRModel)
    result = DataFrame()
    result["t"] = s.ta
    result["S"] = s.Sa
    result["I"] = s.Ia
    result["R"] = s.Ra
    result
end

end # module

The function, live, contains the logic behind the simulation; each susceptible individual contacts others at rate \( c \); if the contact is infected, then there is a probability \( \beta \) that the susceptible individual will become infected. After infection, individuals recover at rate \( \gamma \).

Some things to note:

• Julia doesn't permit redefining types in the main scope, so if you want to run in the REPL, types can be defined in a module, which can be reloaded.
• Julia doesn't allow subtyping of concrete types. In SimPy, one could define SIRPerson as a subclass of Process. Here, I use members of SIR model to hold instances of SimJulia.Simulation and an array of SimJulia.Process.
• I use an outer constructor method for SIRModel.
• I update S, I, and R every time step, as the output method constructs a DataFrame for the results.
• I've separated the contact rate \( c \) from the probability of infection, \( \beta \), just to clarify their different roles in transmission.
• The exclamation mark in push! denotes that the operation is in place (hence, I should really use activate! etc.).

Here is the model in action.

# Load simulation code and libraries
include("sir_sim.jl")
using SimJulia
using Gadfly

# Set parameters
# All Float64 so use decimal points
beta = 0.1
c = 1.0
gamma = 0.05

# Set initial conditions
# All Int64
S0 = 99
I0 = 1
R0 = 0
N0 = S0+I0+R0

# Initialise and run
sirmodel = SIRSim.SIRModel(N0,beta,c,gamma,S0,I0,R0);
SIRSim.activate(sirmodel);
SIRSim.run(sirmodel,1000.0);
# Collate results
result=SIRSim.out(sirmodel)

# Plot using Gadfly
p=plot(results,x="t",y="I",Geom.line)

This all seems a lot more complex than the stochastic model using Gillespie's algorithm I wrote about before. However, this approach makes it easier to build more complex models:

• The distribution of passage times (here the recovery time) can easily be changed to a distribution other than an exponential.
• Infectivity terms and contact rates can be allowed to vary in a complex way across individuals.
• Rather than assuming a well-mixed population, we can easily assume a contact network, such that individuals only contact a subset of the population.

Differential equation modeling with Julia

In the last few examples of epidemiological modeling, I focused on stochastic models. However, most published models are in the form of a set of differential equations. In R, I often use simecol. The ODE package in Julia offers some functionality for solving ordinary differential equations, which should be very familiar to users of Matlab/Octave.

Here's the old favourite, the susceptible-infected-recovered model again.

# Load libraries
using ODE
using DataFrames
using Gadfly

# Define the model
# Has to return a column vector
function SIR(t,x)
    S=x[1]
    I=x[2]
    R=x[3]
    beta=x[4]
    gamma=x[5]
    N=S+I
    dS=-beta*S*I/N
    dI=beta*S*I/N-gamma*I
    dR=gamma*I
    return([dS;dI;dR;0;0])
end

# Initialise model
t = linspace(0,500,101);
inits=[9999,1,0,0.1,0.05];

# Run model
result=ode23(SIR,t,inits);

# Collate results in DataFrame
df=DataFrame();
df["t"]=result[1];
df["S"]=result[2][:,1];
df["I"]=result[2][:,2];
df["R"]=result[2][:,3];

# Plot using Gadfly
p=Gadfly.plot(df,x="t",y="I",Geom.line)

There are a few things to note; firstly, the syntax of the ode23 command expects only the time and a state vector. For parameter values, one can hard-code them in the model description (not good) or pass them as additional variables (with gradient 0, so they don't change). I don't really like either option. Perhaps one can use macros to do much the same job, but I'm not that well acquainted with Julia yet.

An SIR model in Julia

I've been playing around a bit with a relatively new language, Julia, that is dynamic, like R and Python, but purports to have performance not dissimilar to C (see here for my setup). My initial attempts at coding up a simple SIR type epidemiological model, like this were disappointing, which I subsequently discovered was due to how I was storing the results. Here is my second try; you'll need to use an up-to-date build of Julia from the git repository, especially in order to get plotting working.

using DataFrames
using Distributions

function sir(beta,gamma,N,S0,I0,R0,tf)
    t = 0
    S = S0
    I = I0
    R = R0
    ta=DataFrames.DataArray(Float64,0)
    Sa=DataFrames.DataArray(Float64,0)
    Ia=DataFrames.DataArray(Float64,0)
    Ra=DataFrames.DataArray(Float64,0)
    while t < tf
        push!(ta,t)
        push!(Sa,S)
        push!(Ia,I)
        push!(Ra,R)
        pf1 = beta*S*I
        pf2 = gamma*I
        pf = pf1+pf2
        dt = rand(Distributions.Exponential(1/pf))
        t = t+dt
        if t>tf
            break
        end
        ru = rand()
        if ru<(pf1/pf)
            S=S-1
            I=I+1
        else
            I=I-1
            R=R+1
        end
    end
    results = DataFrames.DataFrame()
    results["t"] = ta
    results["S"] = Sa
    results["I"] = Ia
    results["R"] = Ra
    return(results)
end

We can now run the model as follows, with parameter values $\beta$ = 0.1/10000 and $\gamma$ = 0.05, initial conditions $S(0)$ =9999, $I(0)$=1, $R(0)$=0, and a simulation time of 1000.

s=sir(0.1/10000,0.05,10000,9999,1,0,1000)

We can plot using Winston, and save to a PNG file.

using Winston
wp = plot(s[:,"t"],s[:,"I"])
setattr(wp,"xlabel","t")
setattr(wp,"ylabel","I")
file(wp,"sir_winston.png")

Comparing the performance of R and Rcpp for a simple epidemiological model

While R is great for statistics, like many dynamic languages, it's terribly slow in loops. This can make dynamical models, like those commonly used in epidemiology, rather slow. However, thanks to the wonderful Rcpp package, such simulations can be sped up considerably. Let's take the 'standard' susceptible-infected-recovered (SIR) model. A deterministic version of this model is as follows.

\[\begin{align} \frac{dS(t)}{dt} & = -\beta S(t) I(t) \cr \frac{dI(t)}{dt} & = \beta S(t) I(t)- \gamma I(t) \cr \frac{dR(t)}{dt} & = \gamma I(t) \end{align}\]
Let's consider a stochastic version of the SIR model.

\[ \begin{align} {\rm Transition} & \quad {\rm Rate} \cr S \rightarrow S-1,\; I \rightarrow I+1 & \quad \beta S(t) I(t) \cr I \rightarrow I-1,\; R \rightarrow R+1 & \quad \gamma I(t) \end{align}\]
We can simulate this model using the Gillespie method, also known as the stochastic simulation algorithm, or SSA. In plain R, the stochastic model is as follows.

sir <- function(beta, gamma, N, S0, I0, R0, tf) {
    time <- 0
    S <- S0
    I <- I0
    R <- R0
    ta <- numeric(0)
    Sa <- numeric(0)
    Ia <- numeric(0)
    Ra <- numeric(0)
    while (time < tf) {
        ta <- c(ta, time)
        Sa <- c(Sa, S)
        Ia <- c(Ia, I)
        Ra <- c(Ra, R)
        pf1 <- beta * S * I
        pf2 <- gamma * I
        pf <- pf1 + pf2
        dt <- rexp(1, rate = pf)
        time <- time + dt
        if (time > tf) {
            break
        }
        ru <- runif(1)
        if (ru < (pf1/pf)) {
            S <- S - 1
            I <- I + 1
        } else {
            I <- I - 1
            R <- R + 1
        }
        if (I == 0) {
            break
        }
    }
    results <- data.frame(time = ta, S = Sa, I = Ia, R = Ra)
    # print(head(results))
    return(results)
}

Here is the version written using Rcpp, inlining the C++ code.

library(Rcpp)
cppFunction('
  List sirc(double beta, double gamma, double N, double S0, double I0, double R0, double tf){
    double t = 0;
    double S = S0;
    double I = I0;
    double R = R0;
    std::vector<double> ta;
    std::vector<double> Sa;
    std::vector<double> Ia;
    std::vector<double> Ra;
    do{
      ta.push_back(t);
      Sa.push_back(S);
      Ia.push_back(I);
      Ra.push_back(R);
      double pf1 = beta*S*I;
      double pf2 = gamma*I;
      double pf = pf1+pf2;
      double dt = rexp(1,pf)[0];
      t += dt;
      double r = runif(1)[0];
      if(r<pf1/pf){
        S--;
        I++;
      }else{
        I--;
        R++;
      }
      if(I==0){break;}
    } while (t<=tf && (I>0));
  return List::create(_["time"] = ta, _["S"] = Sa, _["I"] = Ia, _["R"]=Ra);
  }'
)

Now we run both the vanilla R and Rcpp versions of the same model 100 times, with parameter values \( \beta \) = 0.1/10000 and \( \gamma \) = 0.05, initial conditions \( S(0) \) =9999, \( I(0) \)=1, \( R(0) \)=0, and a simulation time of 1000.

N <- 100  # Run 100 times
set.seed(4)
sirsim <- sum(replicate(N, system.time(x <- sir(0.1/10000, 0.05, 10000, 9999, 
    1, 0, 1000))["elapsed"]), trim = 0.05)
set.seed(4)
sircsim <- sum(replicate(N, system.time(x <- sirc(0.1/10000, 0.05, 10000, 9999, 
    1, 0, 1000))["elapsed"]), trim = 0.05)

As you can see, the Rcpp version is considerably faster.

paste("Plain R:", sirsim)

## [1] "Plain R: 114.178"

paste("Rcpp version:", sircsim)

## [1] "Rcpp version: 0.175000000000028"

paste("R took", sirsim/sircsim, "times longer")

## [1] "R took 652.445714285608 times longer"