Add new interface for symbolic indexing

- DiffEqArray now stores the new `SymbolCache` struct, which defines and
  implements its interface to query symbols
- DiffEqArray supports symbolically indexing parameters, if provided

Author: AayushSabharwal

src/tabletraits.jl

@@ -7,12 +7,12 @@ function Tables.rows(A::AbstractDiffEqArray)
         N = length(A.u[1])
         names = [
             :timestamp,
-            (A.syms !== nothing ? (A.syms[i] for i in 1:N) :
+            (A.sc !== nothing && A.sc.syms !== nothing ? (A.sc.syms[i] for i in 1:N) :
              (Symbol("value", i) for i in 1:N))...,
         ]
         types = Type[eltype(A.t), (eltype(A.u[1]) for _ in 1:N)...]
     else
-        names = [:timestamp, A.syms !== nothing ? A.syms[1] : :value]
+        names = [:timestamp, A.sc !== nothing && A.sc.syms !== nothing ? A.sc.syms[1] : :value]
         types = Type[eltype(A.t), VT]
     end
     return AbstractDiffEqArrayRows(names, types, A.t, A.u)

src/vector_of_array.jl

@@ -33,6 +33,27 @@ mutable struct VectorOfArray{T, N, A} <: AbstractVectorOfArray{T, N, A}
 end
 # VectorOfArray with an added series for time
 
+
+struct SymbolCache{S,T,U}
+  syms::S
+  indepsym::T
+  paramsyms::U
+end
+
+is_indep_sym(sc::SymbolCache, sym) = isequal(sc.indepsym, sym)
+is_indep_sym(::SymbolCache{S,Nothing}, _) where {S} = false
+state_sym_to_index(sc::SymbolCache, sym) = findfirst(isequal(sym), sc.syms)
+state_sym_to_index(::SymbolCache{Nothing}, _) = nothing
+is_state_sym(sc::SymbolCache, sym) = !isnothing(state_sym_to_index(sc, sym))
+param_sym_to_index(sc::SymbolCache, sym) = findfirst(isequal(sym), sc.paramsyms)
+param_sym_to_index(::SymbolCache{S,T,Nothing}, _) where {S,T} = nothing
+is_param_sym(sc::SymbolCache, sym) = !isnothing(param_sym_to_index(sc, sym))
+
+Base.copy(VA::SymbolCache) = typeof(VA)(
+    (VA.syms===nothing) ? nothing : copy(VA.syms),
+    (VA.indepsym===nothing) ? nothing : copy(VA.indepsym),
+    (VA.paramsyms===nothing) ? nothing : copy(VA.paramsyms),
+  )
 """
 ```julia
 DiffEqArray(u::AbstractVector,t::AbstractVector)
@@ -53,11 +74,10 @@ A[1,:]  # all time periods for f(t)
 A.t
 ```
 """
-mutable struct DiffEqArray{T, N, A, B, C, D, E, F} <: AbstractDiffEqArray{T, N, A}
+mutable struct DiffEqArray{T, N, A, B, C, E, F} <: AbstractDiffEqArray{T, N, A}
   u::A # A <: AbstractVector{<: AbstractArray{T, N - 1}}
   t::B
-  syms::C
-  indepsym::D
+  sc::C
   observed::E
   p::F
 end
@@ -94,11 +114,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}, syms=nothing, indepsym=nothing, observed=nothing, p=nothing) where {T, N} = DiffEqArray{eltype(T), N, typeof(vec), typeof(ts), typeof(syms), typeof(indepsym), typeof(observed), typeof(p)}(vec, ts, syms, indepsym, observed, p)
+DiffEqArray(vec::AbstractVector{T}, ts, ::NTuple{N}, syms=nothing, indepsym=nothing, observed=nothing, p=nothing) where {T, N} = DiffEqArray{eltype(T), N, typeof(vec), typeof(ts), SymbolCache{typeof(syms), typeof(indepsym), Nothing}, typeof(observed), typeof(p)}(vec, ts, SymbolCache(syms, indepsym, nothing), observed, p)
 # Assume that the first element is representative of all other elements
 DiffEqArray(vec::AbstractVector,ts::AbstractVector, syms=nothing, indepsym=nothing, observed=nothing, p=nothing) = DiffEqArray(vec, ts, (size(vec[1])..., length(vec)), syms, indepsym, observed, p)
 function DiffEqArray(vec::AbstractVector{VT},ts::AbstractVector, syms=nothing, indepsym=nothing, observed=nothing, p=nothing) 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)
+  DiffEqArray{T, N+1, typeof(vec), typeof(ts), SymbolCache{typeof(syms), typeof(indepsym), Nothing}, typeof(observed), typeof(p)}(vec, ts, SymbolCache(syms, indepsym, nothing), observed, p)
 end
 
 # Interface for the linear indexing. This is just a view of the underlying nested structure
@@ -138,37 +158,39 @@ Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, i::Int,::Co
 Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, ::Colon,i::Int) where {T, N} = A.u[i]
 Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, i::Int,II::AbstractArray{Int}) where {T, N} = [A.u[j][i] for j in II]
 Base.@propagate_inbounds function Base.getindex(A::AbstractDiffEqArray{T, N},sym) where {T, N}
-  if issymbollike(sym) && A.syms !== nothing
-    i = findfirst(isequal(Symbol(sym)),A.syms)
-  else
-    i = sym
-  end
-
-  if i === nothing
-    if issymbollike(sym) && A.indepsym !== nothing && Symbol(sym) == A.indepsym
-      A.t
+  if issymbollike(sym) && !isnothing(A.sc)
+    if is_indep_sym(A.sc, sym)
+      return A.t
+    elseif is_state_sym(A.sc, sym)
+      return getindex.(A.u, state_sym_to_index(A.sc, sym))
+    elseif is_param_sym(A.sc, sym)
+      return A.p[param_sym_to_index(A.sc, sym)]
+    else
+      return observed(A, sym, :)
+    end
+  elseif all(issymbollike, sym) && !isnothing(A.sc)
+    if all(Base.Fix1(is_param_sym, A.sc), sym)
+      return getindex.((A,), sym)
     else
-      observed(A,sym,:)
+      return [getindex.((A,), sym, i) for i in eachindex(A.t)]
     end
   else
-    Base.getindex.(A.u, i)
+    return getindex.(A.u, sym)
   end
 end
 Base.@propagate_inbounds function Base.getindex(A::AbstractDiffEqArray{T, N},sym,args...) where {T, N}
-  if issymbollike(sym) && A.syms !== nothing
-    i = findfirst(isequal(Symbol(sym)),A.syms)
-  else
-    i = sym
-  end
-
-  if i === nothing
-    if issymbollike(sym) && A.indepsym !== nothing && Symbol(sym) == A.indepsym
-      A.t[args...]
+  if issymbollike(sym) && !isnothing(A.sc)
+    if is_indep_sym(A.sc, sym)
+      return A.t[args...]
+    elseif is_state_sym(A.sc, sym)
+      return A[sym][args...]
     else
-      observed(A,sym,args...)
+      return observed(A, sym, args...)
     end
+  elseif all(issymbollike, sym) && !isnothing(A.sc)
+    return reduce(vcat, map(s -> A[s, args...]', sym))
   else
-    Base.getindex.(A.u, i, args...)
+    return getindex.(A.u, sym)
   end
 end
 Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, I::Int...) where {T, N} = A.u[I[end]][Base.front(I)...]
@@ -230,8 +252,7 @@ tuples(VA::DiffEqArray) = tuple.(VA.t,VA.u)
 Base.copy(VA::AbstractDiffEqArray) = typeof(VA)(
     copy(VA.u),
     copy(VA.t),
-    (VA.syms===nothing) ? nothing : copy(VA.syms),
-    (VA.indepsym===nothing) ? nothing : copy(VA.indepsym),
+    (VA.sc===nothing) ? nothing : copy(VA.sc),
     (VA.observed===nothing) ? nothing : copy(VA.observed),
     (VA.p===nothing) ? nothing : copy(VA.p)
   )