labelledarrays.jl

using ModelingToolkit, StaticArrays, LinearAlgebra, LabelledArrays
using DiffEqBase, ForwardDiff
using Test

# Define some variables
@parameters t σ ρ β
@variables x(t) y(t) z(t)
D = Differential(t)

# Define a differential equation
eqs = [D(x) ~ σ * (y - x),
    D(y) ~ t * x * (ρ - z) - y,
    D(z) ~ x * y - β * z]

@named de = ODESystem(eqs)
ff = ODEFunction(de, [x, y, z], [σ, ρ, β], jac = true)

a = @SVector [1.0, 2.0, 3.0]
b = SLVector(x = 1.0, y = 2.0, z = 3.0)
c = [1.0, 2.0, 3.0]
p = SLVector(σ = 10.0, ρ = 26.0, β = 8 / 3)
@test ff(a, p, 0.0) isa SVector
@test typeof(ff(b, p, 0.0)) <: SLArray
@test ff(c, p, 0.0) isa Vector
@test ff(a, p, 0.0) == ff(b, p, 0.0)
@test ff(a, p, 0.0) == ff(c, p, 0.0)

@test ff.jac(a, p, 0.0) isa SMatrix
@test typeof(ff.jac(b, p, 0.0)) <: SMatrix
@test ff.jac(c, p, 0.0) isa Matrix
@test ff.jac(a, p, 0.0) == ff.jac(b, p, 0.0)
@test ff.jac(a, p, 0.0) == ff.jac(c, p, 0.0)

# Test similar_type
@test ff(b, p, ForwardDiff.Dual(0.0, 1.0)) isa SLArray
d = LVector(x = 1.0, y = 2.0, z = 3.0)
@test ff(d, p, ForwardDiff.Dual(0.0, 1.0)) isa LArray
@test ff.jac(b, p, ForwardDiff.Dual(0.0, 1.0)) isa SArray
@test eltype(ff.jac(b, p, ForwardDiff.Dual(0.0, 1.0))) <: ForwardDiff.Dual
@test ff.jac(d, p, ForwardDiff.Dual(0.0, 1.0)) isa Array
@inferred ff.jac(d, p, ForwardDiff.Dual(0.0, 1.0))
@test eltype(ff.jac(d, p, ForwardDiff.Dual(0.0, 1.0))) <: ForwardDiff.Dual