Merge pull request #78 from SciML/revert-77-revert-66-sy/interp_elim_var_access

Revert "Revert "Add symbol based indexing for interpolated solutions""

Author: ChrisRackauckas

src/interpolation.jl

@@ -248,7 +248,9 @@ times t (sorted), with values u and derivatives ks
   end
 end
 
-
+@inline function interpolant(Θ,id::AbstractDiffEqInterpolation,dt,y₀,y₁,dy₀,dy₁,idxs,::Type{Val{D}}) where D
+  error("$(string(typeof(id))) for $(D)th order not implemented")
+end
 ##################### Hermite Interpolants
 
 """

src/solutions/ode_solutions.jl

@@ -15,9 +15,53 @@ struct ODESolution{T,N,uType,uType2,DType,tType,rateType,P,A,IType,DE} <: Abstra
   destats::DE
   retcode::Symbol
 end
-(sol::ODESolution)(t,deriv::Type=Val{0};idxs=nothing,continuity=:left) = sol.interp(t,idxs,deriv,sol.prob.p,continuity)
+(sol::ODESolution)(t,deriv::Type=Val{0};idxs=nothing,continuity=:left) = sol(t,deriv,idxs,continuity)
 (sol::ODESolution)(v,t,deriv::Type=Val{0};idxs=nothing,continuity=:left) = sol.interp(v,t,idxs,deriv,sol.prob.p,continuity)
 
+function (sol::ODESolution)(t::Real,deriv,idxs::Nothing,continuity)
+  sol.interp(t,idxs,deriv,sol.prob.p,continuity)
+end
+
+function (sol::ODESolution)(t::AbstractVector{<:Real},deriv,idxs::Nothing,continuity)
+  augment(sol.interp(t,idxs,deriv,sol.prob.p,continuity), sol)
+end
+
+function (sol::ODESolution)(t::Real,deriv,idxs::Integer,continuity)
+  sol.interp(t,idxs,deriv,sol.prob.p,continuity)
+end
+function (sol::ODESolution)(t::Real,deriv,idxs::AbstractVector{<:Integer},continuity)
+  sol.interp(t,idxs,deriv,sol.prob.p,continuity)
+end
+function (sol::ODESolution)(t::AbstractVector{<:Real},deriv,idxs::Integer,continuity)
+  sol.interp(t,idxs,deriv,sol.prob.p,continuity)
+end
+function (sol::ODESolution)(t::AbstractVector{<:Real},deriv,idxs::AbstractVector{<:Integer},continuity)
+  sol.interp(t,idxs,deriv,sol.prob.p,continuity)
+end
+
+function (sol::ODESolution)(t::Real,deriv,idxs,continuity)
+  issymbollike(idxs) || error("Incorrect specification of `idxs`")
+  augment(sol.interp([t],nothing,deriv,sol.prob.p,continuity), sol)[idxs][1]
+end
+
+function (sol::ODESolution)(t::Real,deriv,idxs::AbstractVector,continuity)
+  all(issymbollike.(idxs)) || error("Incorrect specification of `idxs`")
+  interp_sol = augment(sol.interp([t],nothing,deriv,sol.prob.p,continuity), sol)
+  [first(interp_sol[idx]) for idx in idxs]
+end
+
+function (sol::ODESolution)(t::AbstractVector{<:Real},deriv,idxs,continuity)
+  issymbollike(idxs) || error("Incorrect specification of `idxs`")
+  interp_sol = augment(sol.interp(t,nothing,deriv,sol.prob.p,continuity), sol)
+  DiffEqArray(interp_sol[idxs], t)
+end
+
+function (sol::ODESolution)(t::AbstractVector{<:Real},deriv,idxs::AbstractVector,continuity)
+  all(issymbollike.(idxs)) || error("Incorrect specification of `idxs`")
+  interp_sol = augment(sol.interp(t,nothing,deriv,sol.prob.p,continuity), sol)
+  DiffEqArray([[interp_sol[idx][i] for idx in idxs] for i in 1:length(t)], t)
+end
+
 function build_solution(
         prob::Union{AbstractODEProblem,AbstractDDEProblem},
         alg,t,u;timeseries_errors=length(u)>2,

src/solutions/solution_interface.jl

@@ -14,6 +14,12 @@ Base.size(A::AbstractNoTimeSolution) = size(A.u)
 
 Base.show(io::IO, m::MIME"text/plain", A::AbstractNoTimeSolution) = (print(io,"u: ");show(io,m,A.u))
 
+# For augmenting system information to enable symbol based indexing of interpolated solutions
+function augment(A::DiffEqArray, sol::AbstractODESolution)
+  observed = has_observed(sol.prob.f) ? sol.prob.f.observed : DEFAULT_OBSERVED
+  DiffEqArray(A.u, A.t, sol.prob.f.syms,getindepsym(sol),observed,sol.prob.p)
+end
+
 # Symbol Handling
 
 # For handling ambiguities
@@ -294,13 +300,15 @@ end
 
 sym_to_index(sym,sol::SciMLSolution) = sym_to_index(sym,getsyms(sol))
 sym_to_index(sym,syms) = findfirst(isequal(Symbol(sym)),syms)
-issymbollike(x) = x isa Symbol ||
+function 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
+                  Symbol(parameterless_type(typeof(x))) == :Operation || Symbol(parameterless_type(typeof(x))) == Symbol("Symbolics.Operation") ||
+                  Symbol(parameterless_type(typeof(x))) == :Variable || Symbol(parameterless_type(typeof(x))) == Symbol("Symbolics.Variable") ||
+                  Symbol(parameterless_type(typeof(x))) == :Sym || Symbol(parameterless_type(typeof(x))) == Symbol("Symbolics.Sym") ||
+                  Symbol(parameterless_type(typeof(x))) == :Num || Symbol(parameterless_type(typeof(x))) == Symbol("Symbolics.Num") ||
+                  Symbol(parameterless_type(typeof(x))) == :Term || Symbol(parameterless_type(typeof(x))) == Symbol("Symbolics.Term")       
+end
 
 function diffeq_to_arrays(sol,plot_analytic,denseplot,plotdensity,tspan,axis_safety,vars,int_vars,tscale,strs)
   if tspan === nothing

test/downstream/observables.jl

@@ -0,0 +1,121 @@
+using ModelingToolkit, OrdinaryDiffEq, Test
+
+@parameters t σ ρ β
+@variables x(t) y(t) z(t)
+D = Differential(t)
+
+eqs = [D(x) ~ σ*(y-x),
+       D(y) ~ x*(ρ-z)-y,
+       D(z) ~ x*y - β*z]
+
+lorenz1 = ODESystem(eqs,name=:lorenz1)
+lorenz2 = ODESystem(eqs,name=:lorenz2)
+
+@parameters γ
+@variables a(t), α(t)
+connections = [0 ~ lorenz1.x + lorenz2.y + a*γ,
+               α ~ 2lorenz1.x + a*γ]
+sys = ODESystem(connections,t,[a,α],[γ],systems=[lorenz1,lorenz2])
+sys_simplified = structural_simplify(sys)
+
+u0 = [lorenz1.x => 1.0,
+      lorenz1.y => 0.0,
+      lorenz1.z => 0.0,
+      lorenz2.x => 0.0,
+      lorenz2.y => 1.0,
+      lorenz2.z => 0.0,
+      a => 2.0]
+
+p  = [lorenz1.σ => 10.0,
+      lorenz1.ρ => 28.0,
+      lorenz1.β => 8/3,
+      lorenz2.σ => 10.0,
+      lorenz2.ρ => 28.0,
+      lorenz2.β => 8/3,
+      γ => 2.0]
+
+tspan = (0.0,100.0)
+prob = ODEProblem(sys_simplified,u0,tspan,p)
+sol = solve(prob,Rodas4())
+
+@test sol[lorenz1.x] isa Vector
+@test sol[lorenz1.x,2] isa Float64
+@test sol[lorenz1.x,:] isa Vector
+@test length(sol[lorenz1.x,1:5]) == 5
+@test sol[α] isa Vector
+@test sol[α,3] isa Float64
+@test length(sol[α,5:10]) == 6
+
+# Check if indexing using variable names from interpolated solution works
+interpolated_sol = sol(0.0:1.0:10.0)
+@test interpolated_sol[α] isa Vector
+@test interpolated_sol[α,:] isa Vector
+@test interpolated_sol[α,2] isa Float64
+@test length(interpolated_sol[α,1:5]) == 5
+@test interpolated_sol[α] ≈ 2interpolated_sol[lorenz1.x] .+ interpolated_sol[a].*2.0
+
+
+sol1 = sol(0.0:1.0:10.0)
+@test sol1.u isa Vector
+@test first(sol1.u) isa Vector
+@test length(sol1.u) == 11
+@test length(sol1.t) == 11
+
+sol2 = sol(0.1)
+@test sol2 isa Vector
+@test length(sol2) == length(states(sys_simplified))
+@test first(sol2) isa Real
+
+sol3 = sol(0.0:1.0:10.0, idxs=[lorenz1.x, lorenz2.x])
+@test sol3.u isa Vector
+@test first(sol3.u) isa Vector
+@test length(sol3.u) == 11
+@test length(sol3.t) == 11
+@test_throws ErrorException sol(0.0:1.0:10.0, idxs=[lorenz1.x, 1])
+
+sol4 = sol(0.1, idxs=[lorenz1.x, lorenz2.x])
+@test sol4 isa Vector
+@test length(sol4) == 2
+@test first(sol4) isa Real
+@test_throws ErrorException sol(0.1, idxs=[lorenz1.x, 1])
+
+sol5 = sol(0.0:1.0:10.0, idxs=lorenz1.x)
+@test sol5.u isa Vector
+@test first(sol5.u) isa Real
+@test length(sol5.u) == 11
+@test length(sol5.t) == 11
+@test_throws ErrorException sol(0.0:1.0:10.0, idxs=1.2)
+
+sol6 = sol(0.1, idxs=lorenz1.x)
+@test sol6 isa Real
+@test_throws ErrorException sol(0.1, idxs=1.2)
+
+sol7 = sol(0.0:1.0:10.0, idxs=[2,1])
+@test sol7.u isa Vector
+@test first(sol7.u) isa Vector
+@test length(sol7.u) == 11
+@test length(sol7.t) == 11
+
+sol8 = sol(0.1, idxs=[2,1])
+@test sol8 isa Vector
+@test length(sol8) == 2
+@test first(sol8) isa Real
+
+sol9 = sol(0.0:1.0:10.0, idxs=2)
+@test sol9.u isa Vector
+@test first(sol9.u) isa Real
+@test length(sol9.u) == 11
+@test length(sol9.t) == 11
+
+sol10 = sol(0.1, idxs=2)
+@test sol10 isa Real
+
+
+#=
+using Plots
+plot(sol,vars=(lorenz2.x,lorenz2.z))
+plot(sol,vars=(α,lorenz2.z))
+plot(sol,vars=(lorenz2.x,α))
+plot(sol,vars=α)
+plot(sol,vars=(t,α))
+=#

test/downstream/symbol_indexing.jl

@@ -27,5 +27,6 @@ end
 if !is_APPVEYOR && GROUP == "Downstream"
     activate_downstream_env()
     @time @safetestset "Ensembles of Zero Length Solutions" begin include("downstream/ensemble_zero_length.jl") end
+    @time @safetestset "Symbol and integer based indexing of interpolated solutions" begin include("downstream/symbol_indexing.jl") end
 end
 end