Feature Request: ComponentArrays labels in plot recipe

[metanoid]

I notice that in DiffEqFlux#master there is now documented support for ComponentArrays.jl.
I've checked, and ODEProblem can handle both u0 and p being ComponentArray-type with no code changes.

A nice-to-have extra feature: If u0 is a ComponentArray, then by default calling plot(sol) on the solutions should use the names from u0 as the labels.

[ChrisRackauckas]

Yeah, that would be a nice addition.

[jonniedie]

This is mostly available now. Surface-level components can be plotted with a symbol like vars=:x. More deeply-nested components can be plotted with a string like vars="lotka.x". If the component has subcomponents, it will plot all of them as well (with appropriate labels for each, of course). Standard index arguments like vars=1, vars=[1, 3], vars=(1, 2, 3), or vars=(f, 1, 2) will automatically label the plot with component names too. Some of the features aren't there yet, such as tuple arguments like vars=(f, :x, :y) or even vars=(:x, :y, :z). Here is a concrete example:

using ComponentArrays
using DifferentialEquations
using Parameters: @unpack
using Plots


## Equations
# Lorenz system
function lorenz!(D, u, p, t; f=0.0)
    @unpack σ, ρ, β = p
    @unpack x, y, z = u
    
    D.x = σ*(y - x)
    D.y = x*(ρ - z) - y - f
    D.z = x*y - β*z
    return nothing
end

# Lotka-Volterra system
function lotka!(D, u, p, t; f=0.0)
    @unpack α, β, γ, δ = p
    @unpack x, y = u
    
    D.x =  α*x - β*x*y + f
    D.y = -γ*y + δ*x*y
    return nothing
end

# Composed Lorenz and Lotka-Volterra system
function composed!(D, u, p, t)
    c = p.c #coupling parameter
    @unpack lorenz, lotka = u
    
    lorenz!(D.lorenz, lorenz, p.lorenz, t, f=c*lotka.x)
    lotka!(D.lotka, lotka, p.lotka, t, f=c*lorenz.x)
    return nothing
end


## Inputs
# Simulation timespan
tspan = (0.0, 20.0)

# Parameters
lorenz_p = (σ=10.0, ρ=28.0, β=8/3)
lotka_p = (α=2/3, β=4/3, γ=1.0, δ=1.0)
comp_p = (lorenz=lorenz_p, lotka=lotka_p, c=0.01)

# Initial conditions
lorenz_ic = ComponentArray(x=0.0, y=0.0, z=0.0)
lotka_ic = ComponentArray(x=1.0, y=1.0)
comp_ic = ComponentArray(lorenz=lorenz_ic, lotka=lotka_ic)

# Make problem and solve
sol = ODEProblem(composed!, comp_ic, tspan, comp_p) |> solve


## Plotting
# Plotting by number will automatically produce labels
p1 = plot(sol; vars=(1, 2, 3))

# Plotting by component name will include all of its subcomponents
p2 = plot(sol; vars=:lotka)

# Nested components can be accessed through strings
p3 = plot(sol; vars=["lotka.x", "lorenz.x"])

plot(p1, plot(p2, p3; layout=(2,1)); layout=(1,2))

[ChrisRackauckas]

How did it hook in? There shouldn't be any extra steps to support those things, so I think it might've been done incorrectly.

[jonniedie]

Here's how I hooked into it. Any DiffEqBase.AbstractODESolution{T,N,C} where C<:AbstractVector{<:ComponentArray} gets string labels for the legend unless syms are already supplied by the user.

I did have to commit type piracy in one place (which I just noticed is inconsistent w.r.t cleansyms/cleansym naming):

DiffEqBase.cleansyms(syms::AbstractArray{<:String}) = DiffEqBase.cleansyms.(syms)
DiffEqBase.cleansyms(syms::String) = syms

I can go back on this, though, if you think it can create issues.

[jonniedie]

Since cleansyms converts ModelingToolkit symnames to a string and then replaces ₊ with ., I added a fall-through method for strings that are already in that form.

[ChrisRackauckas]

That sounds about right. Odd it takes so much though: MTK only needs like one or two tie-in points, and so does LabelledArrays

[jonniedie]

Yeah, I'll have to dive back in to see why I needed that extra hook into interpret_vars. There was something going on where if I tried to do vars=[:lotka₊x], it would plot time vs. time.