Solution handling for ComponentArrays: how to extract a variable by name?

[scheidan]

ComponentArrays.jl is very helpful to handle states and parameters and works well with DifferentialEquations (e.g. plots are automatically labeled #617).

Unfortunaly, it seem we lack a good interface to extract variables from solutions objects. Something like sol(t, idxs=:x) or sol(t, idxs=[:x, :y]) would feel very natural but is currently not implemented:

using OrdinaryDiffEq
using ComponentArrays

# ODE system where u0 is a ComponentVector
function lorenz!(du, u, p, t)
    du.x = p.a * (u.y - u.x)
    du.y = u.x * (p.b - u.z) - u.y
    du.z = u.x * u.y - p.c * u.z
end

u0 = ComponentVector(x=1, y=0, z=0)
p = ComponentVector(a=10.0, b=28.0, c=88.3)
tspan = (0.0, 1.0)
prob = ODEProblem(lorenz!, u0, tspan, p)
sol = solve(prob, Tsit5());
t = 0:0.1:0.3

# --- How to extract solution for :x?
sol(t, idxs=:x)   # ERROR: Indexing symbol x is unknown.
sol(t, idxs=[:x]) # ERROR: Indexing symbol x is unknown.

# --- Workarounds, not exactly pretty...
Array(sol(t))[:x,:]
sol(t, idxs=ComponentArrays.label2index(first(sol.u), "x"))

# --- Interestingly, this works!
sol2 = solve(prob, Tsit5(), save_idxs=:x);
sol2(t)

[paulxshen]

Doesn't sol(t).x work already?

[scheidan]

That would be a nice interface. Unfortunately it does not work:

julia> sol(t).x
ERROR: type DiffEqArray has no field x
Stacktrace:
 [1] getproperty(x::RecursiveArrayTools.DiffEqArray{…}, f::Symbol)
   @ Base ./Base.jl:37
 [2] top-level scope
   @ REPL[15]:1
Some type information was truncated. Use `show(err)` to see complete types.

[ChrisRackauckas]

We would need to forward the getproperty on DiffEqArray. That does seem possible. @AayushSabharwal

[AayushSabharwal]

How about a ComponentArraysExt which implements SII methods for AbstractDiffEqArrays containing componentarrays to make symbolic indexing possible? Forwarding properties to enable sol.x has downsides:

1. It breaks if someone defines a component array with u as one of its symbols
2. It is not in line with how the rest of the DSL works

The extension allows us to support sol[:x] syntax which is exactly how symbolic indexing works right now

[ChrisRackauckas]

👍 That sounds like a better option.