Much better support for labelledarrays

Automatic plot labelled, support with DAE solvers (no allocs!), etc. They should just work everywhere.

Author: ChrisRackauckas

Project.toml

@@ -13,6 +13,7 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 IteratorInterfaceExtensions = "82899510-4779-5014-852e-03e436cf321d"
+LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"

src/DiffEqBase.jl

@@ -7,7 +7,7 @@ using RecipesBase, RecursiveArrayTools,
 
 import Logging, LoggingExtras, TerminalLoggers, ConsoleProgressMonitor, ProgressLogging
 
-import ZygoteRules, ChainRulesCore
+import ZygoteRules, ChainRulesCore, LabelledArrays
 
 using Roots # callbacks
 

src/init.jl

@@ -2,6 +2,7 @@ value(x) = x
 cuify(x) = error("To use LinSolveGPUFactorize, you must do `using CuArrays`")
 promote_u0(u0,p,t0) = u0
 promote_tspan(u0,p,tspan,prob,kwargs) = tspan
+get_tmp(x) = nothing
 
 if VERSION < v"1.4.0-DEV.635"
   # Piracy, should get upstreamed
@@ -61,12 +62,21 @@ function __init__()
     end
 
     function DiffCache(u::AbstractArray{T}, siz, ::Type{Val{chunk_size}}) where {T, chunk_size}
-        DiffCache(u, zeros(ForwardDiff.Dual{nothing,T,chunk_size}, siz...))
+        x = ArrayInterface.restructure(u,zeros(ForwardDiff.Dual{nothing,T,chunk_size}, siz...))
+        DiffCache(u, x)
     end
 
     dualcache(u::AbstractArray, N=Val{ForwardDiff.pickchunksize(length(u))}) = DiffCache(u, size(u), N)
 
-    get_tmp(dc::DiffCache, u::AbstractArray{T}) where T<:ForwardDiff.Dual = reinterpret(T, dc.dual_du)
+    function get_tmp(dc::DiffCache, u::AbstractArray{T}) where T<:ForwardDiff.Dual
+      x = reinterpret(T, dc.dual_du)
+    end
+
+    function DiffEqBase.get_tmp(dc::DiffEqBase.DiffCache, u::LabelledArrays.LArray{T,N,D,Syms}) where {T,N,D,Syms}
+      x = reinterpret(T, dc.dual_du.__x)
+      LArray{T,N,D,Syms}(x)
+    end
+
     get_tmp(dc::DiffCache, u::AbstractArray) = dc.du
 
     bisection(f, tup::Tuple{T,T}, tdir) where {T<:ForwardDiff.Dual} = find_zero(f, tup, Roots.AlefeldPotraShi())

src/solutions/solution_interface.jl

@@ -69,11 +69,12 @@ DEFAULT_PLOT_FUNC(x,y,z) = (x,y,z) # For v0.5.2 bug
                    tspan = nothing, axis_safety = 0.1,
                    vars=nothing)
 
-  int_vars = interpret_vars(vars,sol)
+  syms = getsyms(sol)
+  int_vars = interpret_vars(vars,sol,syms)
   tscale = get(plotattributes, :xscale, :identity)
   plot_vecs,labels = diffeq_to_arrays(sol,plot_analytic,denseplot,
                                       plotdensity,tspan,axis_safety,
-                                      vars,int_vars,tscale)
+                                      vars,int_vars,tscale,syms)
 
   tdir = sign(sol.t[end]-sol.t[1])
   xflip --> tdir < 0
@@ -146,7 +147,17 @@ DEFAULT_PLOT_FUNC(x,y,z) = (x,y,z) # For v0.5.2 bug
   (plot_vecs...,)
 end
 
-function diffeq_to_arrays(sol,plot_analytic,denseplot,plotdensity,tspan,axis_safety,vars,int_vars,tscale)
+function getsyms(sol)
+  if has_syms(sol.prob.f)
+    return sol.prob.f.syms
+  elseif typeof(sol.u[1]) <: Union{LabelledArrays.LArray,LabelledArrays.SLArray}
+    return LabelledArrays.symnames(typeof(sol.u[1]))
+  else
+    return nothing
+  end
+end
+
+function diffeq_to_arrays(sol,plot_analytic,denseplot,plotdensity,tspan,axis_safety,vars,int_vars,tscale,syms)
   if tspan === nothing
     if sol.tslocation == 0
       end_idx = length(sol)
@@ -217,21 +228,21 @@ function diffeq_to_arrays(sol,plot_analytic,denseplot,plotdensity,tspan,axis_saf
     @assert length(var) == dims
   end
   # Should check that all have the same dims!
-  plot_vecs,labels = solplot_vecs_and_labels(dims,int_vars,plot_timeseries,plott,sol,plot_analytic,plot_analytic_timeseries)
+  plot_vecs,labels = solplot_vecs_and_labels(dims,int_vars,plot_timeseries,plott,sol,plot_analytic,plot_analytic_timeseries,syms)
 end
 
-function interpret_vars(vars,sol)
-  if vars !== nothing && has_syms(sol.prob.f)
+function interpret_vars(vars,sol,syms)
+  if vars !== nothing && syms !== nothing
     # Do syms conversion
     tmp_vars = []
     for var in vars
       if typeof(var) <: Symbol
-        var_int = something(findfirst(isequal(var), sol.prob.f.syms), 0)
+        var_int = something(findfirst(isequal(var), syms), 0)
       elseif typeof(var) <: Union{Tuple,AbstractArray} #eltype(var) <: Symbol # Some kind of iterable
         tmp = []
         for x in var
           if typeof(x) <: Symbol
-            push!(tmp,something(findfirst(isequal(x), sol.prob.f.syms), 0))
+            push!(tmp,something(findfirst(isequal(x), syms), 0))
           else
             push!(tmp,x)
           end
@@ -314,14 +325,14 @@ function interpret_vars(vars,sol)
   vars
 end
 
-function add_labels!(labels,x,dims,sol)
+function add_labels!(labels,x,dims,sol,syms)
   lys = []
   for j in 3:dims
     if x[j] == 0
       push!(lys,"t,")
     else
-      if has_syms(sol.prob.f)
-        push!(lys,"$(sol.prob.f.syms[x[j]]),")
+      if syms !== nothing
+        push!(lys,"$(syms[x[j]]),")
       else
         push!(lys,"u$(x[j]),")
       end
@@ -331,8 +342,8 @@ function add_labels!(labels,x,dims,sol)
   if x[2] == 0 && dims == 3
     tmp_lab = "$(lys...)(t)"
   else
-    if has_syms(sol.prob.f) && x[2] != 0
-      tmp = sol.prob.f.syms[x[2]]
+    if syms !== nothing && x[2] != 0
+      tmp = syms[x[2]]
       tmp_lab = "($tmp,$(lys...))"
     else
       if x[2] == 0
@@ -350,14 +361,14 @@ function add_labels!(labels,x,dims,sol)
   labels
 end
 
-function add_analytic_labels!(labels,x,dims,sol)
+function add_analytic_labels!(labels,x,dims,sol,syms)
   lys = []
   for j in 3:dims
     if x[j] == 0 && dims == 3
       push!(lys,"t,")
     else
-      if has_syms(sol.prob.f)
-        push!(lys,string("True ",sol.prob.f.syms[x[j]],","))
+      if syms !== nothing
+        push!(lys,string("True ",syms[x[j]],","))
       else
         push!(lys,"True u$(x[j]),")
       end
@@ -367,8 +378,8 @@ function add_analytic_labels!(labels,x,dims,sol)
   if x[2] == 0
     tmp_lab = "$(lys...)(t)"
   else
-    if has_syms(sol.prob.f)
-      tmp = string("True ",sol.prob.f.syms[x[2]])
+    if syms !== nothing
+      tmp = string("True ",syms[x[2]])
       tmp_lab = "($tmp,$(lys...))"
     else
       tmp_lab = "(True u$(x[2]),$(lys...))"
@@ -396,7 +407,7 @@ function u_n(timeseries::AbstractArray, n::Int,sol,plott,plot_timeseries)
   end
 end
 
-function solplot_vecs_and_labels(dims,vars,plot_timeseries,plott,sol,plot_analytic,plot_analytic_timeseries)
+function solplot_vecs_and_labels(dims,vars,plot_timeseries,plott,sol,plot_analytic,plot_analytic_timeseries,syms)
   plot_vecs = []
   labels = String[]
   for x in vars
@@ -416,7 +427,7 @@ function solplot_vecs_and_labels(dims,vars,plot_timeseries,plott,sol,plot_analyt
       end
       push!(plot_vecs[i],tmp[i])
     end
-    add_labels!(labels,x,dims,sol)
+    add_labels!(labels,x,dims,sol,syms)
   end
 
 
@@ -434,7 +445,7 @@ function solplot_vecs_and_labels(dims,vars,plot_timeseries,plott,sol,plot_analyt
       for i in eachindex(tmp)
         push!(plot_vecs[i],tmp[i])
       end
-      add_analytic_labels!(labels,x,dims,sol)
+      add_analytic_labels!(labels,x,dims,sol,syms)
     end
   end
   plot_vecs = [hcat(x...) for x in plot_vecs]

test/downstream/labelledarrays.jl

@@ -0,0 +1,17 @@
+using OrdinaryDiffEq
+using LabelledArrays
+
+function f(out,du,u,p,t)
+  out.x = - 0.04u.x              + 1e4*u.y*u.z - du.x
+  out.y = + 0.04u.x - 3e7*u.y^2 - 1e4*u.y*u.z - du.y
+  out.z = u.x + u.y + u.z - 1.0
+end
+
+u₀ = LVector(x=1.0, y=0.0, z=0.0)
+du₀ = LVector(x=-0.04, y=0.04, z=0.0)
+tspan = (0.0,100000.0)
+
+differential_vars = LVector(x=true, y=true, z=false)
+prob = DAEProblem(f,du₀,u₀,tspan,differential_vars=differential_vars)
+
+sol = solve(prob, DImplicitEuler())

test/runtests.jl

@@ -40,6 +40,7 @@ if !is_APPVEYOR && GROUP == "Downstream"
     @time @safetestset "Ensemble Analysis" begin include("downstream/ensemble_analysis.jl") end
     @time @safetestset "Default linsolve with structure" begin include("downstream/default_linsolve_structure.jl") end
     @time @safetestset "Callback Merging Tests" begin include("downstream/callback_merging.jl") end
+    @time @safetestset "LabelledArrays Tests" begin include("downstream/labelledarrays.jl") end
     @time @safetestset "ODE Event Tests" begin include("downstream/ode_event_tests.jl") end
     @time @safetestset "Event Detection Tests" begin include("downstream/event_detection_tests.jl") end
     @time @safetestset "PSOS and Energy Conservation Event Detection" begin include("downstream/psos_and_energy_conservation.jl") end