How to solve a combination of ODEs and PDEs and organize the solution

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()

odetest

1 Like