Enable symbol based indexing of interpolated solutions by adding extra fields to DiffEqArray type

Author: sharanry

Project.toml

@@ -26,11 +26,12 @@ julia = "1.3"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ForwardDiff", "NLsolve", "OrdinaryDiffEq", "Test", "Unitful", "Random", "StructArrays", "Zygote"]
+test = ["ForwardDiff", "NLsolve", "OrdinaryDiffEq", "Pkg", "Test", "Unitful", "Random", "StructArrays", "Zygote"]

src/RecursiveArrayTools.jl

@@ -18,11 +18,13 @@ using DocStringExtensions
   include("zygote.jl")
 
   export VectorOfArray, DiffEqArray, AbstractVectorOfArray, AbstractDiffEqArray,
-         vecarr_to_arr, vecarr_to_vectors, tuples
+       AllObserved, vecarr_to_arr, vecarr_to_vectors, tuples
 
   export recursivecopy, recursivecopy!, vecvecapply, copyat_or_push!,
          vecvec_to_mat, recursive_one, recursive_mean, recursive_bottom_eltype,
-         recursive_unitless_bottom_eltype, recursive_unitless_eltype
+         recursive_unitless_bottom_eltype, recursive_unitless_eltype, parameterless_type,
+         issymbollike
+
 
   export ArrayPartition
 

src/vector_of_array.jl

@@ -3,9 +3,13 @@ mutable struct VectorOfArray{T, N, A} <: AbstractVectorOfArray{T, N, A}
   u::A # A <: AbstractVector{<: AbstractArray{T, N - 1}}
 end
 # VectorOfArray with an added series for time
-mutable struct DiffEqArray{T, N, A, B} <: AbstractDiffEqArray{T, N, A}
+mutable struct DiffEqArray{T, N, A, B, C, D, E, F} <: AbstractDiffEqArray{T, N, A}
   u::A # A <: AbstractVector{<: AbstractArray{T, N - 1}}
   t::B
+  syms::C
+  indepsym::D
+  observed::E
+  p::F
 end
 
 Base.Array(VA::AbstractVectorOfArray{T,N,A}) where {T,N,A <: AbstractVector{<:AbstractVector}} = reduce(hcat,VA.u)
@@ -20,10 +24,11 @@ VectorOfArray(vec::AbstractVector{T}, ::NTuple{N}) where {T, N} = VectorOfArray{
 VectorOfArray(vec::AbstractVector) = VectorOfArray(vec, (size(vec[1])..., length(vec)))
 VectorOfArray(vec::AbstractVector{VT}) where {T, N, VT<:AbstractArray{T, N}} = VectorOfArray{T, N+1, typeof(vec)}(vec)
 
-DiffEqArray(vec::AbstractVector{T}, ts, ::NTuple{N}) where {T, N} = DiffEqArray{eltype(T), N, typeof(vec), typeof(ts)}(vec, ts)
+DiffEqArray(vec::AbstractVector{T}, ts, ::NTuple{N}) where {T, N} = DiffEqArray{eltype(T), N, typeof(vec), typeof(ts), Nothing, Nothing, Nothing, Nothing}(vec, ts, nothing, nothing, nothing, nothing)
 # Assume that the first element is representative of all other elements
 DiffEqArray(vec::AbstractVector,ts::AbstractVector) = DiffEqArray(vec, ts, (size(vec[1])..., length(vec)))
-DiffEqArray(vec::AbstractVector{VT},ts::AbstractVector) where {T, N, VT<:AbstractArray{T, N}} = DiffEqArray{T, N+1, typeof(vec), typeof(ts)}(vec, ts)
+DiffEqArray(vec::AbstractVector{VT},ts::AbstractVector) where {T, N, VT<:AbstractArray{T, N}} = DiffEqArray{T, N+1, typeof(vec), typeof(ts), Nothing, Nothing, Nothing, Nothing}(vec, ts, nothing, nothing, nothing, nothing)
+DiffEqArray(vec::AbstractVector{VT},ts::AbstractVector, syms::Vector{Symbol}, indepsym::Symbol, observed::Function, p) where {T, N, VT<:AbstractArray{T, N}} = DiffEqArray{T, N+1, typeof(vec), typeof(ts), typeof(syms), typeof(indepsym), typeof(observed), typeof(p)}(vec, ts, syms, indepsym, observed, p)
 
 # Interface for the linear indexing. This is just a view of the underlying nested structure
 @inline Base.firstindex(VA::AbstractVectorOfArray) = firstindex(VA.u)
@@ -38,6 +43,72 @@ Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, I::Int)
 Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, I::Colon) where {T, N} = VA.u[I]
 Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, I::AbstractArray{Int}) where {T, N} = VectorOfArray(VA.u[I])
 Base.@propagate_inbounds Base.getindex(VA::AbstractDiffEqArray{T, N}, I::AbstractArray{Int}) where {T, N} = DiffEqArray(VA.u[I],VA.t[I])
+# 
+Base.@pure __parameterless_type(T) = Base.typename(T).wrapper
+parameterless_type(x) = parameterless_type(typeof(x))
+parameterless_type(x::Type) = __parameterless_type(x)
+
+### Abstract Interface
+struct AllObserved
+end
+issymbollike(x) = x isa Symbol ||
+                  x isa AllObserved ||
+                  Symbol(parameterless_type(typeof(x))) == :Operation ||
+                  Symbol(parameterless_type(typeof(x))) == :Variable ||
+                  Symbol(parameterless_type(typeof(x))) == :Sym ||
+                  Symbol(parameterless_type(typeof(x))) == :Num ||
+                  Symbol(parameterless_type(typeof(x))) == :Term
+
+Base.@propagate_inbounds function Base.getindex(A::AbstractDiffEqArray{T, N},sym) where {T, N}
+  if issymbollike(sym)
+    i = findfirst(isequal(Symbol(sym)),A.syms)
+  else
+    i = sym
+  end
+
+  if i === nothing
+    # TODO: Check if system actually has a indepsym
+    if issymbollike(i) && Symbol(i) == A.indepsym
+      A.t
+    else
+      observed(A,sym,:)
+    end
+  else
+    A[i,:]
+  end
+end
+Base.@propagate_inbounds function Base.getindex(A::AbstractDiffEqArray{T, N},sym,args...) where {T, N}
+  if issymbollike(sym)
+    i = findfirst(isequal(Symbol(sym)),A.syms)
+  else
+    i = sym
+  end
+
+  if i === nothing
+    # TODO: Check if system actually has a indepsym
+    if issymbollike(i) && Symbol(i) == A.indepsym
+      A.t[args...]
+    else
+      observed(A,sym,args...)
+    end
+  else
+    A[i,args...]
+  end
+end
+Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N} where {T, N}, i::Int64, ::Colon) = Base.getindex.(A.u, i)
+
+function observed(A::AbstractDiffEqArray{T, N},sym,i::Int) where {T, N}
+    A.observed(sym,A.u[i],A.p,A.t[i])
+end
+
+function observed(A::AbstractDiffEqArray{T, N},sym,i::AbstractArray{Int}) where {T, N}
+    A.observed.((sym,),A.u[i],(A.p,),A.t[i])
+end
+
+function observed(A::AbstractDiffEqArray{T, N},sym,::Colon) where {T, N}
+    A.observed.((sym,),A.u,(A.p,),A.t)
+end
+# 
 Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, i::Int,::Colon) where {T, N} = [VA.u[j][i] for j in 1:length(VA)]
 Base.@propagate_inbounds function Base.getindex(VA::AbstractVectorOfArray{T,N}, ii::CartesianIndex) where {T, N}
     ti = Tuple(ii)

test/downstream/Project.toml

@@ -0,0 +1,3 @@
+[deps]
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"

test/downstream/symbol_indexing.jl

@@ -0,0 +1,23 @@
+using RecursiveArrayTools, ModelingToolkit, OrdinaryDiffEq, Test
+
+@variables t x(t)  # independent and dependent variables
+@parameters τ       # parameters
+D = Differential(t) # define an operator for the differentiation w.r.t. time
+@variables RHS(t)
+@named fol_separate = ODESystem([ RHS  ~ (1 - x)/τ,
+                                  D(x) ~ RHS ])
+fol_simplified = structural_simplify(fol_separate)                                  
+
+prob = ODEProblem(fol_simplified, [x => 0.0], (0.0,10.0), [τ => 3.0])
+sol = solve(prob, Tsit5())
+
+sol_new = DiffEqArray(
+    sol.u[1:10],
+    sol.t[1:10],
+    sol.prob.f.syms,
+    sol.prob.f.indepsym,
+    sol.prob.f.observed,
+    sol.prob.p
+)
+
+@test sol_new[RHS] ≈ (1 .- sol_new[x])./3.0

test/runtests.jl

@@ -1,6 +1,16 @@
+using Pkg
 using RecursiveArrayTools
 using Test
 
+const GROUP = get(ENV, "GROUP", "All")
+const is_APPVEYOR = ( Sys.iswindows() && haskey(ENV,"APPVEYOR") )
+
+function activate_downstream_env()
+  Pkg.activate("downstream")
+  Pkg.develop(PackageSpec(path=dirname(@__DIR__)))
+  Pkg.instantiate()
+end
+
 @time begin
   @time @testset "Utils Tests" begin include("utils_test.jl") end
   @time @testset "Partitions Tests" begin include("partitions_test.jl") end
@@ -10,4 +20,9 @@ using Test
   @time @testset "Linear Algebra Tests" begin include("linalg.jl") end
   @time @testset "Upstream Tests" begin include("upstream.jl") end
   @time @testset "Adjoint Tests" begin include("adjoints.jl") end
+
+  if !is_APPVEYOR && GROUP == "Downstream"
+    activate_downstream_env()
+    @time @testset "DiffEqArray Indexing Tests" begin include("downstream/symbol_indexing.jl") end
+  end
 end