Porting SIR ODE Model to Julia Part 1

A while ago, I wrote a post about fitting 3-state SIR models in Stan and this is my attempt to port it from R/Stan to Julia/Turing so I can get a feel for the language and to do some basic benchmarks. I am writing this from the perspective of an R programmer who has never used Julia before.

I will restate the model we are trying to reproduce here.

$$\begin{align*}
\frac{\text{d}S}{\text{d}t} &= -\beta \frac{I(t)}{N} S(t) \\
\frac{\text{d}I}{\text{d}t} &= \beta \frac{I(t)}{N} S(t) - \gamma I(t) \\
\frac{\text{d}R}{\text{d}t} &= \gamma I(t)
\end{align*}$$

To start, I will define the ODE system in Julia and do some forward simulations.

using Plots
using LaTeXStrings
using OrdinaryDiffEq
using Distributions
using Random
using Turing

Random.seed!(1234)

function sir!(du, u, p, t)
    S, I, R = u
    b, g = p 
    du[1] = -b * I * S         # dS/dt
    du[2] =  b * I * S - g * I # dI/dt
    du[3] =  g * I             # dR/dt
end

‍

The exclamation point at the end of the function name is a convention (not a rule) in Julia that alerts users that the function will modify some of its arguments in place, using pass-by-reference semantics. Since $du$ will be modified many times, this is very efficient. This is not possible in R, as R will copy using the copy-on-modify rule. In the first two lines, the arguments are copied from vectors to scalars, which are used in the ODE.

The next block of code sets up the data for simulation.

‍

n = 100
n_pop = 763
pars = (b = 2.0, g = 0.5)
u0 = [(n_pop - 1)/n_pop, 1/n_pop, 0]  # initial state vector
times = range(1, stop = 15, length = n)
tspan = (first(times), last(times))

‍

The pars line creates a named tuple, and tspan will be needed later to set up the ODEProblem, which does not have an equivalent in R/Stan. In order to examine the structure of a Julia object, one can use typeof, which is similar to R's function with the same name, and dump, which is similar to str.

julia> typeof(tspan)
Tuple{Float64, Float64}

julia> dump(tspan)
Tuple{Float64, Float64}
  1: Float64 1.0
  2: Float64 15.0

‍

prob = ODEProblem(sir!, u0, tspan, pars)
sol  = solve(prob, DP5(), saveat = times)

‍

ODEProblem packages all the information necessary for the solver in one convenient object. This object can be passed to different solvers using different algorithms. Julia is big on composability, and these types of patterns repeat often. In our case, we will be using the DP5 solver, which implements the same algorithm as Stan's ode_rk45. Both use Dormand–Prince 5(4) Runge–Kutta method with adaptive step size and are meant for non-stiff systems. One thing to note is that ODEProblem sets the time interval and anchors the initial condition at $t_0$, while solve takes times to evaluate the solution at the requested time points.

Once the solution is generated, we can take a look at some plots.

‍

plot(sol, xlabel = "Days", 
          ylabel = "Population fraction", 
          label  = ["S" "I" "R"],
          title  = "SIR Model with " * L"\beta = 2,\ \gamma = 0.5")

‍

The plot function is quite general and can take recipes (instructions on how to draw), this time from SciML, Julia's Scientific Machine Learning library. This means I don't have to unpack the sol object; I can just pass it in.

Next, we will add some Poisson noise to the solution and show the results.

‍

I = sol[2, :] * n_pop
y = rand.(Poisson.(I))
p = scatter(sol.t, y, label="observed (noisy)", alpha = 0.5)
p = plot!(sol.t, I, label="true")

‍

‍

The first line extracts the second row of the sol matrix, which would be sol[2, ] in R. The second line is where we meet Julia's broadcasting system, invoked by the period character. By default, most Julia functions operate on scalars, so if we tried Poisson(I), Julia would complain with ERROR: MethodError: no method matching Poisson(::Vector{Int64}). The call to Poisson(.) creates a distribution object. This is similar to R's distributional package and the corresponding dist_*() functions. rand performs the sampling and requires the period operator for the same reason as above. You can apply the period operator to your own functions (e.g., my_fun.()) just as you do using built-in functions.

The plotting function calls demonstrate the generality of the "update-in-place" paradigm -- the plot! function adds to the scatter object that had been previously created.

Now that everything is in place, we can define the probabilistic model using Turing.jl.

@model function sir_model(y, n_pop, prob, times)
    b  ~ LogNormal(0, 1)
    g  ~ LogNormal(0, 1)
    S0 ~ Beta(1, 1)
 
    u0 = [S0, 1.0 - S0, 0.0]   # [S(0), I(0), R(0)]
    p  = [b, g]

    sol = solve(
        prob, DP5();
        u0     = u0,
        p      = p,
        saveat = times,
        abstol = 1e-8,
        reltol = 1e-8
    )

    if SciMLBase.successful_retcode(sol) == false
        Turing.@addlogprob!(-Inf)
        return
    end

    I_frac = sol[2, :]
    λ = max.(I_frac .* n_pop, 1e-9)
    
    for i in eachindex(y)
        y[i] ~ Poisson(λ[i])
    end

    return (; λ, I_frac, R0 = b/g)
end

‍

The first thing to note is the @model prefix, which indicates that what follows is a macro invocation. This macro rewrites the function definition into a probabilistic program in Julia. The transformation is handled at compile time and is hidden from the user, in much the same way that Stan hides its C++ code generation. While the syntax looks like ordinary Julia, the semantics differ from those of deterministic programs: statements describe probabilistic relationships and contribute to a joint probability distribution rather than executing as standard assignments.

The syntax is mostly self-explanatory, with a few caveats. SciMLBase.successful_retcode(sol) == false catches ODE integration failures and was necessary to make this model run. When it triggers, we add -Inf to the log-probability (i.e., zero probability), so that the proposal is rejected and the sampler continues to the next iteration. The second guard is to make sure that $\lambda$ doesn't get too small.

In Turing, there’s no separate block for generated quantities like in Stan. If you want per-draw derived values (e.g., λ, I_frac, R0) without rerunning the ODE solver later, return them from the model as a named tuple: return (; λ, I_frac, R0 = b/g).

Finally, we package the model and call the sampler.

‍

model = sir_model(y, n_pop, prob, times)
sampler = NUTS();
n_samples = 500      
n_adapts  = 500     
n_chains  = 4
chains = sample(
    model,      
    sampler,     
    MCMCThreads(),
    n_samples,
    n_chains,
    n_adapts = n_adapts
)

‍

MCMCThreads is telling the sampler to run each chain in a separate thread. This is different from CmdStanR, where each chain runs in a separate OS process.

Fit summary is obtained by running describe(chains).

‍

julia> describe(chains)
Chains MCMC chain (500×17×4 Array{Float64, 3}):

Iterations        = 251:1:750
Number of chains  = 4
Samples per chain = 500
Wall duration     = 4.08 seconds
Compute duration  = 11.11 seconds
parameters        = b, g, S0
internals         = n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, step_size, nom_step_size, lp, logprior, loglikelihood

Summary Statistics

  parameters      mean       std      mcse    ess_bulk    ess_tail      rhat   ess_per_sec 
      Symbol   Float64   Float64   Float64     Float64     Float64   Float64       Float64 

           b    2.0129    0.0243    0.0010    653.6191    828.1861    1.0048       58.8528
           g    0.5081    0.0047    0.0001   1104.0545   1004.9696    1.0052       99.4106
          S0    0.9988    0.0001    0.0000    651.2514    727.3439    1.0023       58.6396


Quantiles

  parameters      2.5%     25.0%     50.0%     75.0%     97.5% 
      Symbol   Float64   Float64   Float64   Float64   Float64 

           b    1.9666    1.9962    2.0128    2.0299    2.0607
           g    0.4986    0.5050    0.5081    0.5113    0.5171
          S0    0.9985    0.9987    0.9988    0.9988    0.9990

‍

Computation of generated quantities, functions of parameters computed post-inference, is done with generated_quantities function.

gq = generated_quantities(sir_model(y, n_pop, prob, times), chains)
R0 = [s.R0 for s in gq] # 500x4 matrix 
R0 = map(s -> s.R0, gq) # redundant, same as above
mean_R0 = mean(s.R0 for s in gq)
R0 = vec(R0)            # 2,000 element vector

‍

The function returns gq, which is a vector of named tuples --- one per posterior draw, 2,000 in total. The second line uses Julia's comprehensions to iterate over gq and pulls out $R_0$. The same thing can be done using the map function, which is similar to R's map in the purrr package. The fourth line is using Julia's (lazy) generator, which can be passed to summary functions like mean. Here mean_R0 $\approx 3.96$, which is close to our preset value of 4. The last line flattens the 500x4 matrix into a 2,000-element vector.

The other two generated quantities (λ and I_frac) will contain 100 (e.g., length(gq[1].λ) = 100) predictions for each draw, since we set up the simulation for 100 time points.

Before we turn to plotting the infection curve predictions, a few words on access elements in the chains object returned by sample.

‍

b_samp  = vec(Array(chains[:b]))
g_samp  = vec(Array(chains[:g]))
S0_samp = vec(Array(chains[:S0]))

scatter(b_samp, g_samp;
        alpha = 0.3,
        xlabel = "β",
        ylabel = "γ",
        title = "Posterior samples of β and γ",
        legend = false)
scatter!([pars.b], [pars.g], ms = 5)

‍

‍

Using the same methods as described above, what's left is to extract 100 random draws from I_frac, put them on the absolute count scale, and plot.

ndraws = min(100, nsamples)
gq_flat = vec(gq)
idx = rand(eachindex(gq), n)
I_pred = [gq_flat[j].I_frac .* n_pop for j in idx]
plot(
    times,
    I_pred,
    label = "",
    alpha = 0.05,
    color = :blue
)
scatter!(times, 
         y,
         label = "",
         color = "red",
         alpha = 0.3)
xlabel!("Days")
ylabel!("Infected (count)")

‍

In the next post, I will run some comparison benchmarks.

Thanks to Nikolas Siccha for reading over the draft and providing helpful comments.

‍