Closed jlperla closed 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
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.