Wednesday, June 19, 2013

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")

No comments:

Post a Comment