I am making a series of compartmental ODE infection models using DifferentialEquations. I would like to be able to plot the sum of two or more compartments (e.g. I might have people with undiagnosed infection and diagnosed infection, and want to plot the total number with infection).
A simple example is that I’d like to sum the entire population each time I change the model to check that I haven’t made a mistake that causes the population size to change.
I’ve put a simple example below that is very unattractive. I can work at changing saveat to make the plot smoother, but I’m hoping there’s something within DifferentialEquations that I’ve missed that will perform this more elegantly. Thanks.
using DifferentialEquations, Plots, StaticArrays
function sir(u, p, t)
S, I, R = u
β, γ = p
N = S + I + R
dS = -β * S * I / N
dI = β * S * I / N - γ * I
dR = γ * I
return @SVector [dS, dI, dR]
end
u₀ = @SVector [.999, .001, 0]
p = [.3, .1]
tspan = (0., 100.)
prob = ODEProblem(sir, u₀, tspan, p)
sol = solve(prob)
function plot_total_size(solution)
y::Array = sum.(solution.u')'
x = solution.t
plot(x, y)
end
plot_total_size(sol)
