Thanks @contradict and @ChrisRackauckas for the suggestion to use ComponentArrays. I have implemented it in my code and it works well.
I did find it non-trival to access the parts of the DiffEq solution using the component labels. Ultimately, I found the label2index()
function which is part of ComponentArrays can be used to find the appropriate indices which can be used with the provided solution analysis tools or plot recipes in DifferentialEquations.jl.
If there is a way to access the components without using label2index()
, please share!
Below is a MWE that highlights how this can work.
using ComponentArrays
using OrdinaryDiffEq
using Parameters: @unpack
using Plots
# Define function to get index of
# each labeled component in solution
idx(sol,var) = label2index(sol.u[1],var)
function odetest()
# Define and solve
# ODE dXo/dt = Xo and
# PDE d(Xp=A,B,C)/dt = Xp
# ---------------------
# Define right-hand-side of ODEs
function dsol_dt!(dsol, sol, p, t)
@unpack Xo,Xp = sol
dsol.Xo = Xo
dsol.Xp = Xp
return nothing
end
# Initial conditions
Xo₀ = 1.0 #
Xp₀ = ComponentArray(A = 1.0, B = 2.0, C = 3.0)
sol₀ = ComponentArray(Xo=Xo₀,Xp=Xp₀)
# Setup problem and solve
tspan = (0.,1.0)
prob = ODEProblem(dsol_dt!,sol₀,tspan)
sol = solve(prob,Tsit5(),)
# Access final solution
# ---------------------
# Access components
Xo = sol[idx(sol,"Xo"),:]
Xp = sol[idx(sol,"Xp"),:]
# Access components with interpolation
t = 0.0:0.1:1.0
Xo_dense = sol(t,idxs = idx(sol,"Xo"))
Xp_dense = sol(t,idxs = idx(sol,"Xp"))
# Plot solution and components
p1 = plot(sol ,title="Full solution")
p2 = plot(sol,idxs = idx(sol,"Xo"),title="Only Xo")
p3 = plot(sol,idxs = idx(sol,"Xp"),title="Only Xp")
p = plot(p1,p2,p3,layout=(1,3))
display(p)
end
odetest()