# How to vectorize a SIR Model to add age-stratification with the DifferentialEquations.jl package?

**URL:** <https://discourse.julialang.org/t/how-to-vectorize-a-sir-model-to-add-age-stratification-with-the-differentialequations-jl-package/49375>\
**Category:** Modelling & Simulations\
**Tags:** question, diffeq\
**Created:** [October 31, 2020, 3:58pm UTC](https://discourse.julialang.org/t/how-to-vectorize-a-sir-model-to-add-age-stratification-with-the-differentialequations-jl-package/49375 "2020-10-31T15:58:12Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![PietroMonticone](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pietromonticone/32/44881_2.png) [@PietroMonticone](https://discourse.julialang.org/u/PietroMonticone)\
**Post date:** [October 31, 2020, 3:58pm UTC](https://discourse.julialang.org/t/how-to-vectorize-a-sir-model-to-add-age-stratification-with-the-differentialequations-jl-package/49375/1 "2020-10-31T15:58:13Z")

</div>

Consider the following simple SIR model I’ve written in the DifferentialEquations.jl framework:

```julia
# Import 
using DifferentialEquations

# Parameter
𝒫 = [0.1,0.2] # Model Parameters
ℬ = [0.8,0.1, 0.0] # Initial condition
𝒯 = (0.0,365.0) # Time 

# Model
function Φ!(du,u,p,t)
    S,I,R = u
    β,γ = p
    
    du[1] = -β*S*I
    du[2] = β*S*I - γ*I
    du[3] = γ*I
end

# Problem Definition
problem = ODEProblem(Φ!, ℬ, 𝒯, 𝒫)

# Problem Solution
solution = solve(problem);

```

I would like to extend it with age-stratification in a way that’s both scalable (at least up to 20 age groups per compartment) and compact (vector form).

In the following code chunk I try to define an uncoupled age-specific SIR with only 3 age groups:

```julia
# Parameter
𝒫 = [0.14,0.12,0.2,0.2] # Model Parameters
ℬ = [0.9,0.9,0.9,0.1,0.1,0.1,0.0,0.0,0.0] # Initial condition
𝒯 = (0.0,365.0) # Time 

# Model
function Φ!(du,u,p,t)
    s1,s2,s3,i1,i2,i3,r1,r2,r3 = u 
    β1,β2,β3,γ = p
    
    du[1] = -β1*s1*i1
    du[2] = -β2*s2*i2
    du[3] = -β3*s3*i3
    
    du[4] = β1*s1*i1-γ*i1
    du[5] = β2*s2*i2-γ*i2
    du[6] = β3*s3*i3-γ*i3

    du[7] = γ*i1
    du[8] = γ*i2
    du[9] = γ*i3
end

# Problem Definition
problem = ODEProblem(Φ!, ℬ, 𝒯, 𝒫)

# Problem Solution
solution = solve(problem);

```

This works fine but it’s neither compact nor scalable. Can you help me vectorize it in a way to make it so?

---

<div class="post-metadata">

**Author:** ![baggepinnen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/baggepinnen/32/693_2.png) [@baggepinnen](https://discourse.julialang.org/u/baggepinnen)\
**Post date:** [October 31, 2020, 4:45pm UTC](https://discourse.julialang.org/t/how-to-vectorize-a-sir-model-to-add-age-stratification-with-the-differentialequations-jl-package/49375/2 "2020-10-31T16:45:19Z")

</div>

You could use the type `Particles` from [MonteCarloMeasurements.jl](https://github.com/baggepinnen/MonteCarloMeasurements.jl) which is made for doing vectorized operations. On your problem, it would look like this, where the different model parameters are given as a vector to `Particles`

```julia
using OrdinaryDiffEq, MonteCarloMeasurements

# Parameter
𝒫 = [Particles([0.14, 0.12]), Particles([0.2, 0.2])] # Model Parameters
ℬ = eltype(𝒫).([0.9, 0.1, 0]) # Initial condition
𝒯 = (0.0,365.0) # Time 

# Model
function Φ!(du,u,p,t)
    S,I,R = u
    β,γ = p
    
    du[1] = -β*S*I
    du[2] = β*S*I - γ*I
    du[3] = γ*I
end

# Problem Definition
problem = ODEProblem(Φ!, ℬ, 𝒯, 𝒫)

# Problem Solution
solution = solve(problem, Tsit5());
mcplot(solution.t, Matrix(Array(solution)'))

```

 ![sir](https://global.discourse-cdn.com/julialang/original/3X/8/9/89cdfdc809579c397f70bf49c2a42d2028065964.png)

To make sure initial conditions always sum to one when there are distributions of initial conditions, it’s convenient to caluclate one from the others, e.g.,

```julia
S = Particles(...)
I = Particles(...)
R = 1 - S - I

```

To simulate something like an uncertainty in the initial conditions, you could do something like

```julia
𝒫 = [0.1 0.2] # Model Parameters
S = 0.9 ± 0.005 # Initial condition
R = 0
I = 1 - S - R
ℬ = [S, I, R]
𝒯 = (0.0,365.0) # Time 

```

where the `± (\pm)` operator creates 2000 normally distributed `Particles`.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [October 31, 2020, 5:30pm UTC](https://discourse.julialang.org/t/how-to-vectorize-a-sir-model-to-add-age-stratification-with-the-differentialequations-jl-package/49375/3 "2020-10-31T17:30:30Z")

</div>

Just make your initial condition a matrix and then:

```julia
function Φ!(du,u,p,t)
    S = @view u[:,1]
    I = @view u[:,2]
    R = @view u[:,3]
    dS = @view du[:,1]
    dI = @view du[:,2]
    dR = @view du[:,3]
    β = @view p[1:3]
    γ = p[4]
    
    @. dS = -β*S*I
    @. dI = β*S*I-γ*I
    @. dR = γ*I
end

```

You can then use LabelledArrays.jl to make this style even simpler:

[https://github.com/SciML/LabelledArrays.jl](https://github.com/SciML/LabelledArrays.jl)

---

<div class="post-metadata">

**Author:** ![eliassno](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/eliassno/32/18917_2.png) [@eliassno](https://discourse.julialang.org/u/eliassno)\
**Post date:** [October 31, 2020, 6:07pm UTC](https://discourse.julialang.org/t/how-to-vectorize-a-sir-model-to-add-age-stratification-with-the-differentialequations-jl-package/49375/4 "2020-10-31T18:07:31Z")

</div>

I was considering the same thing as Chris.  
[https://diffeq.sciml.ai/stable/tutorials/ode\_example/#ode\_other\_types](https://diffeq.sciml.ai/stable/tutorials/ode_example/#ode_other_types)

The solutions are equivalent once you match the indices of the initial conditions and the views that Chris proposes, e.g. by setting the initial conditions as:

```julia
ℬ = [0.9 0.9 0.9; 0.1 0.1 0.1; 0.0 0.0 0.0]' # Initial condition

```

* * *

For _scalability_ (perhaps for compactess as well), one could store the parameters as

```julia
𝒫 = [[0.14,0.12,0.2],0.2] # Model Parameters

```

and avoid hard-coding the number of age-groups in `Φ!`

```julia
function Φ!(du,u,p,t)
    S = @view u[:,1]
    I = @view u[:,2]
    R = @view u[:,3]
    dS = @view du[:,1]
    dI = @view du[:,2]
    dR = @view du[:,3]
    β, γ = p

    @. du[:,1] = -β*S*I
    @. du[:,2] = β*S*I - γ*I
    @. du[:,3] = γ*I
end

```

This does have a significant performance impact when calling `solve`, so perhaps it’s better to count the number of age groups.

```julia
N = size(u,1)
β = @view p[1:N]
γ = p[N+1]

```