QuantEcon / lecture-source-jl

Source files for "Lectures in Quantitative Economics" -- Julia version
https://julia.quantecon.org
BSD 3-Clause "New" or "Revised" License
131 stars 53 forks source link

Port the covid lecture to julia #828

Closed jlperla closed 4 years ago

jlperla commented 4 years ago

https://python.quantecon.org/sir_model.html

This would use the DIffEq features. At some point in the future it might be near to look at it for SDE or DiffEqFlux examples.

jlperla commented 4 years ago
using StochasticDiffEq
using Random
using SparseArrays
function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    @inbounds begin
        du[1] = -β*c*I/N*S
        du[2] = β*c*I/N*S - γ*I
        du[3] = γ*I
    end
    nothing
end;
# Define a sparse matrix by making a dense matrix and setting some values as not zero
A = zeros(3,2)
A[1,1] = 1
A[2,1] = 1
A[2,2] = 1
A[3,2] = 1
A = SparseArrays.sparse(A);
# Make `g` write the sparse matrix values
function sir_noise!(du,u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    ifrac = β*I/N*S
    rfrac = γ*I
    du[1,1] = -sqrt(ifrac)
    du[2,1] = sqrt(ifrac)
    du[2,2] = -sqrt(rfrac)
    du[3,2] = sqrt(rfrac)
end;
δt = 0.1
tmax = 40.0
tspan = (0.0,tmax)
t = 0.0:δt:tmax;
u0 = [990.0,10.0,0.0]; # S,I,R
p = [0.05,10.0,0.25]; # β,c,γ
prob_sde = SDEProblem(sir_ode!,sir_noise!,u0,tspan,p,noise_rate_prototype=A)

using Plots
plot(sol_sde)
ensemble_prob = EnsembleProblem(prob_sde)
sim = solve(ensemble_prob,SRA1(),EnsembleThreads(),trajectories=100)
summ = EnsembleSummary(sim,quantiles=[0.025,0.975])
plot(summ,labels="Middle 95%")
summ = EnsembleSummary(sim,quantiles=[0.25,0.75])
plot!(summ,labels="Middle 50%", legend=:topleft)
sim = solve(ensemble_prob,SRA1(),EnsembleThreads(),trajectories=100)
summ = EnsembleSummary(sim,quantiles=[0.025,0.975])
plot(summ,labels="Middle 95%")
summ = EnsembleSummary(sim,quantiles=[0.25,0.75])
plot!(summ,labels="Middle 50%", legend=:topleft)

And see https://aip.scitation.org/doi/10.1063/1.481811 and https://docs.sciml.ai/latest/tutorials/sde_example/#Example-4:-Systems-of-SDEs-with-Non-Diagonal-Noise-1 and point people to https://github.com/epirecipes/sir-julia