I finally got an SIR model to work with SIR combined with ModelToolkit. Here is the code:
using Plots using DataFrames using ModelingToolkit using DiffEqBase using DiffEqJump using LightGraphs using SpecialGraphs const nedgess = 3000 const nvertss = 500 const dgs = DiGraph(nvertss, nedgess); # Degree of each node: 2nd argument 5 const dgr = random_regular_digraph(nvertss, 5) function setupSystem(graph) nverts = length(vertices(graph)) nedges = length(edges(graph)) # Allow for different beta on each edgse @parameters t β[1:nedges] γ[1:nverts] ; @variables S[1:nverts](t); @variables I[1:nverts](t); @variables R[1:nverts](t); rxsS = [Reaction(β[i],[I[src(e)], S[dst(e)]], [I[dst(e)]], [1,1], ) for (i,e) ∈ enumerate(edges(graph))] rxsI = [Reaction(γ[v],[I[v]], [R[v]]) # note: src to src, yet there is no edges for v ∈ vertices(graph)] rxs = vcat(rxsS, rxsI); vars = vcat(S,I,R); params = vcat(β,γ); rs = ReactionSystem(rxs, t, vars, params); js = convert(JumpSystem, rs); println("Completed: convert(JumpSystem)") S0 = ones(nverts) I0 = zeros(nverts) R0 = zeros(nverts) S0 = 0.; # One person is infected I0 = 1.; R0 = 1. - S0 - I0 vars0 = vcat(S0, I0, R0); # Two column vectors γ = fill(0.25, nverts); β = fill(0.50, nedges); params = vcat(β,γ) initial_state = [convert(Variable,state) => vars0[i] for (i,state) in enumerate(states(js))]; initial_params = [convert(Variable,par) => params[i] for (i,par) in enumerate(parameters(js))]; tspan = (0.0,20.0) @time dprob = DiscreteProblem(js, initial_state, tspan, initial_params) println("Completed: DiscreteProblem") @time jprob = JumpProblem(js, dprob, NRM()) println("Completed: JumpProblem") @time sol = solve(jprob, SSAStepper()) println("Completed: solve") return sol end function processData(sol) nverts = nvertss nedges = nedgess println("nverts=$nverts, nedges= $nedges") dfs = convert(Matrix, DataFrame(sol)) Sf = dfs[1:nverts,:] If = dfs[nverts+1:2*nverts,:] Rf = dfs[2*nverts+1:3*nverts,:] Savg = (sum(Sf; dims=1)') / nverts Iavg = (sum(If; dims=1)') / nverts Ravg = (sum(Rf; dims=1)') / nverts print(Savg) return Savg, Iavg, Ravg end # Times: sol.t # Solution at nodes: sol.u # sol.u |> length == 120 (3 * nverts) sol = setupSystem(dgr) Savg, Iavg, Ravg = processData(sol) plot(sol.t, Savg) plot(sol.t, Iavg) plot(sol.t, Ravg)
A single run for
tspan=(0.,20.) takes about 18.85 sec, with most of the time taken in
JumpProblem (18.55 sec). The network has 500 nodes with 5 edges per node. There were only 955 discrete times. I do not have a sense on how well my code is optimized, although I do not see how poor optimization on my part could slow down the
JumpProblem routine, although I do not really know.
If anybody has any insight, I would be greatly appreciative. I will try and copy some code from Simon Frost and see if I can duplicate the results. Simon did not use the ModelingKit.