ComponentArrays - Julia Tips

https://github.com/SciML/ComponentArrays.jl are array blocks that can be accessed through a named index. You can compose differential equations without a DSL like ModelingToolkit.

using ComponentArrays
using OrdinaryDiffEq
using SimpleUnPack: @unpack

Composing Lorenz and Lotka-Volterra systems¶

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

lorenz! (generic function with 1 method)

tspan = (0.0, 20.0)
lorenz_p = (σ=10.0, ρ=28.0, β=8/3)
lorenz_ic = ComponentArray(x=1.0, y=0.0, z=0.0)
lorenz_prob = ODEProblem(lorenz!, lorenz_ic, tspan, lorenz_p)

ODEProblem with uType ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(x = 1, y = 2, z = 3)}}} and tType Float64. In-place: true
Non-trivial mass matrix: false
timespan: (0.0, 20.0)
u0: ComponentVector{Float64}(x = 1.0, y = 0.0, z = 0.0)

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

lotka! (generic function with 1 method)

lotka_p = (α=2/3, β=4/3, γ=1.0, δ=1.0)
lotka_ic = ComponentArray(x=1.0, y=1.0)
lotka_prob = ODEProblem(lotka!, lotka_ic, tspan, lotka_p)

ODEProblem with uType ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(x = 1, y = 2)}}} and tType Float64. In-place: true
Non-trivial mass matrix: false
timespan: (0.0, 20.0)
u0: ComponentVector{Float64}(x = 1.0, y = 1.0)

Composed system¶

function composed!(D, u, p, t)
    c = p.c
    @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

composed! (generic function with 1 method)

comp_p = (lorenz=lorenz_p, lotka=lotka_p, c=0.01)
comp_ic = ComponentArray(lorenz=lorenz_ic, lotka=lotka_ic)
comp_prob = ODEProblem(composed!, comp_ic, tspan, comp_p)

ODEProblem with uType ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(lorenz = ViewAxis(1:3, Axis(x = 1, y = 2, z = 3)), lotka = ViewAxis(4:5, Axis(x = 1, y = 2)))}}} and tType Float64. In-place: true
Non-trivial mass matrix: false
timespan: (0.0, 20.0)
u0: ComponentVector{Float64}(lorenz = (x = 1.0, y = 0.0, z = 0.0), lotka = (x = 1.0, y = 1.0))

Solve problem¶

We can solve the composed system

@time comp_sol = solve(comp_prob, Tsit5())

  1.994077 seconds (9.52 M allocations: 487.184 MiB, 3.36% gc time, 99.92% compilation time)

retcode: Success
Interpolation: specialized 4th order "free" interpolation
t: 212-element Vector{Float64}:
  0.0
  6.181923459793355e-5
  0.001559599560883202
  0.006394935946393775
  0.013776712797758741
  0.02347245169786879
  0.03678566059424463
  0.05408354037685106
  0.07651720070071483
  0.10506469374556526
  ⋮
 19.452597867180252
 19.536361940044152
 19.616570634372064
 19.683312330118675
 19.75334863535752
 19.82519153496308
 19.90193886774304
 19.961959843144488
 20.0
u: 212-element Vector{ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(lorenz = ViewAxis(1:3, Axis(x = 1, y = 2, z = 3)), lotka = ViewAxis(4:5, Axis(x = 1, y = 2)))}}}}:
 ComponentVector{Float64}(lorenz = (x = 1.0, y = 0.0, z = 0.0), lotka = (x = 1.0, y = 1.0))
 ComponentVector{Float64}(lorenz = (x = 0.9993825332995702, y = 0.0017297323086565809, z = 5.3446465987803404e-8), lotka = (x = 0.9999594060148846, y = 0.9999999987452488))
 ComponentVector{Float64}(lorenz = (x = 0.9848617238680404, y = 0.04328548286876235, z = 3.345818410264088e-5), lotka = (x = 0.9989762765470666, y = 0.9999992015931267))
 ComponentVector{Float64}(lorenz = (x = 0.9435309277428389, y = 0.17315896258724386, z = 0.0005348624869886599), lotka = (x = 0.995807754416456, y = 0.9999865878570847))
 ComponentVector{Float64}(lorenz = (x = 0.8955393966556862, y = 0.3608576553810656, z = 0.0023190404194028294), lotka = (x = 0.9909871819231854, y = 0.9999378385749176))
 ComponentVector{Float64}(lorenz = (x = 0.8571475696036538, y = 0.5934068634410052, z = 0.006257549534967387), lotka = (x = 0.9846882286627692, y = 0.9998198993251624))
 ComponentVector{Float64}(lorenz = (x = 0.8438308062761597, y = 0.8989388830858138, z = 0.014318940045646641), lotka = (x = 0.9761047980235618, y = 0.9995589075448201))
 ComponentVector{Float64}(lorenz = (x = 0.8848284273360745, y = 1.295639242022747, z = 0.029644098993840833), lotka = (x = 0.9650759928491792, y = 0.9990503279023125))
 ComponentVector{Float64}(lorenz = (x = 1.024333279084553, y = 1.8540009985929193, z = 0.06048521637103717), lotka = (x = 0.9510004246829277, y = 0.9981098074072455))
 ComponentVector{Float64}(lorenz = (x = 1.3383599726505555, y = 2.7196995468525897, z = 0.1298160858161551), lotka = (x = 0.9334978469692644, y = 0.9964644993644551))
 ⋮
 ComponentVector{Float64}(lorenz = (x = -9.064631857530895, y = -14.877252434310533, z = 17.760874103272354), lotka = (x = 4.422635823724514, y = 0.17324199287817596))
 ComponentVector{Float64}(lorenz = (x = -13.811865949936532, y = -17.922696402263533, z = 29.386433501641935), lotka = (x = 4.562965880056151, y = 0.2321683047928164))
 ComponentVector{Float64}(lorenz = (x = -13.702376261466696, y = -8.934642794798085, z = 38.32335367692658), lotka = (x = 4.665350192640776, y = 0.310315608934773))
 ComponentVector{Float64}(lorenz = (x = -9.071373188595333, y = -1.2612587705334184, z = 35.40783242411148), lotka = (x = 4.719512927305387, y = 0.39710824156833563))
 ComponentVector{Float64}(lorenz = (x = -4.352151102602904, y = 0.8867382750389377, z = 29.339360847412355), lotka = (x = 4.735325083566584, y = 0.5157123949509358))
 ComponentVector{Float64}(lorenz = (x = -1.7092895631663705, y = 0.6205504347279385, z = 24.06928249638489), lotka = (x = 4.692035819844095, y = 0.673686133522652))
 ComponentVector{Float64}(lorenz = (x = -0.6233861025237412, y = 0.10995263610682696, z = 19.586618557743876), lotka = (x = 4.559790939002256, y = 0.890399924521133))
 ComponentVector{Float64}(lorenz = (x = -0.3637662402424692, y = -0.16445840977969642, z = 16.69034772760858), lotka = (x = 4.383817086175004, y = 1.0970571894305075))
 ComponentVector{Float64}(lorenz = (x = -0.3257678139414668, y = -0.3131414537101895, z = 15.083239168925475), lotka = (x = 4.237279404979757, y = 1.244397424479974))

Extract a state variable

ts = 0:0.1:20
map(u -> u.lotka.x, comp_sol(ts).u)

201-element Vector{Float64}:
 1.0
 0.9365676092645404
 0.8792888960774066
 0.8317135260247215
 0.7981596571749999
 0.7657768717657337
 0.7244406470610838
 0.6818769363751237
 0.6418734059143623
 0.6050209536382868
 ⋮
 3.8987389826474566
 4.115032177752235
 4.321996549587615
 4.505632713020478
 4.647217815320861
 4.727597748970701
 4.715179098220895
 4.564371224772363
 4.237279404979758

We can also unit test one of the component systems

lotka_sol = solve(lotka_prob, Tsit5())
lotka_sol(1.0).x

0.5742850840860039

Update a subset of the elements in ComponentArrays¶

using ForwardDiff
newval = ForwardDiff.Dual(3.0, 1.0)
v = ComponentVector(a = 1.0, b = 2.0, c = 3.0)

ComponentVector{Float64}(a = 1.0, b = 2.0, c = 3.0)

This will throw an error because the new value has a different type than the existing values

try
    ComponentVector(v; a = newval)
catch e
    println("Error: ", e)
end

Error: MethodError(Float64, (Dual{Nothing}(3.0,1.0),), 0x000000000000989d)

Find the common type of existing values and new value

T = promote_type(eltype(newval), Float64)

ForwardDiff.Dual{Nothing, Float64, 1}

Convert the existing values to the common type

vT = ComponentVector{T}(v)

ComponentVector{ForwardDiff.Dual{Nothing, Float64, 1}}(a = Dual{Nothing}(1.0,0.0), b = Dual{Nothing}(2.0,0.0), c = Dual{Nothing}(3.0,0.0))

Create a new ComponentVector with the updated value

ComponentVector(vT; a = newval)

ComponentVector{ForwardDiff.Dual{Nothing, Float64, 1}}(a = Dual{Nothing}(3.0,1.0), b = Dual{Nothing}(2.0,0.0), c = Dual{Nothing}(3.0,0.0))

This notebook was generated using Literate.jl.