Merge pull request #1988 from AayushSabharwal/syms_rework

Symbolic indexing rework

Author: ChrisRackauckas

Project.toml

@@ -40,6 +40,7 @@ SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
@@ -76,6 +77,7 @@ Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "0.1.0"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
+SymbolicIndexingInterface = "0.1"
 SymbolicUtils = "0.19"
 Symbolics = "4.9"
 UnPack = "0.1, 1.0"

src/ModelingToolkit.jl

@@ -34,6 +34,9 @@ RuntimeGeneratedFunctions.init(@__MODULE__)
 
 using RecursiveArrayTools
 
+import SymbolicIndexingInterface
+import SymbolicIndexingInterface: independent_variables, states, parameters
+export independent_variables, states, parameters
 import SymbolicUtils
 import SymbolicUtils: istree, arguments, operation, similarterm, promote_symtype,
                       Symbolic, Term, Add, Mul, Pow, Sym, FnType,
@@ -96,29 +99,8 @@ abstract type AbstractODESystem <: AbstractTimeDependentSystem end
 abstract type AbstractMultivariateSystem <: AbstractSystem end
 abstract type AbstractOptimizationSystem <: AbstractTimeIndependentSystem end
 
-"""
-$(TYPEDSIGNATURES)
-
-Get the set of independent variables for the given system.
-"""
-function independent_variables end
-
 function independent_variable end
 
-"""
-$(TYPEDSIGNATURES)
-
-Get the set of states for the given system.
-"""
-function states end
-
-"""
-$(TYPEDSIGNATURES)
-
-Get the set of parameters variables for the given system.
-"""
-function parameters end
-
 # this has to be included early to deal with depency issues
 include("structural_transformation/bareiss.jl")
 function complete end
@@ -203,7 +185,7 @@ export Differential, expand_derivatives, @derivatives
 export Equation, ConstrainedEquation
 export Term, Sym
 export SymScope, LocalScope, ParentScope, DelayParentScope, GlobalScope
-export independent_variables, independent_variable, states, parameters, equations, controls,
+export independent_variable, equations, controls,
        observed, structure, full_equations
 export structural_simplify, expand_connections, linearize, linearization_function
 export DiscreteSystem, DiscreteProblem

src/systems/abstractsystem.jl

@@ -148,7 +148,7 @@ function independent_variable(sys::AbstractSystem)
 end
 
 #Treat the result as a vector of symbols always
-function independent_variables(sys::AbstractSystem)
+function SymbolicIndexingInterface.independent_variables(sys::AbstractSystem)
     systype = typeof(sys)
     @warn "Please declare ($systype) as a subtype of `AbstractTimeDependentSystem`, `AbstractTimeIndependentSystem` or `AbstractMultivariateSystem`."
     if isdefined(sys, :iv)
@@ -160,9 +160,13 @@ function independent_variables(sys::AbstractSystem)
     end
 end
 
-independent_variables(sys::AbstractTimeDependentSystem) = [getfield(sys, :iv)]
-independent_variables(sys::AbstractTimeIndependentSystem) = []
-independent_variables(sys::AbstractMultivariateSystem) = getfield(sys, :ivs)
+function SymbolicIndexingInterface.independent_variables(sys::AbstractTimeDependentSystem)
+    [getfield(sys, :iv)]
+end
+SymbolicIndexingInterface.independent_variables(sys::AbstractTimeIndependentSystem) = []
+function SymbolicIndexingInterface.independent_variables(sys::AbstractMultivariateSystem)
+    getfield(sys, :ivs)
+end
 
 iscomplete(sys::AbstractSystem) = isdefined(sys, :complete) && getfield(sys, :complete)
 
@@ -462,15 +466,15 @@ function namespace_expr(O, sys, n = nameof(sys))
     end
 end
 
-function states(sys::AbstractSystem)
+function SymbolicIndexingInterface.states(sys::AbstractSystem)
     sts = get_states(sys)
     systems = get_systems(sys)
     unique(isempty(systems) ?
            sts :
            [sts; reduce(vcat, namespace_variables.(systems))])
 end
 
-function parameters(sys::AbstractSystem)
+function SymbolicIndexingInterface.parameters(sys::AbstractSystem)
     ps = get_ps(sys)
     systems = get_systems(sys)
     unique(isempty(systems) ? ps : [ps; reduce(vcat, namespace_parameters.(systems))])
@@ -508,7 +512,9 @@ end
 states(sys::AbstractSystem, v) = renamespace(sys, v)
 parameters(sys::AbstractSystem, v) = toparam(states(sys, v))
 for f in [:states, :parameters]
-    @eval $f(sys::AbstractSystem, vs::AbstractArray) = map(v -> $f(sys, v), vs)
+    @eval function $f(sys::AbstractSystem, vs::AbstractArray)
+        map(v -> $f(sys, v), vs)
+    end
 end
 
 flatten(sys::AbstractSystem, args...) = sys
@@ -572,6 +578,22 @@ function time_varying_as_func(x, sys::AbstractTimeDependentSystem)
     return x
 end
 
+SymbolicIndexingInterface.is_indep_sym(sys::AbstractSystem, sym) = isequal(sym, get_iv(sys))
+
+function SymbolicIndexingInterface.state_sym_to_index(sys::AbstractSystem, sym)
+    findfirst(isequal(sym), SymbolicIndexingInterface.states(sys))
+end
+function SymbolicIndexingInterface.is_state_sym(sys::AbstractSystem, sym)
+    !isnothing(SymbolicIndexingInterface.state_sym_to_index(sys, sym))
+end
+
+function SymbolicIndexingInterface.param_sym_to_index(sys::AbstractSystem, sym)
+    findfirst(isequal(sym), SymbolicIndexingInterface.parameters(sys))
+end
+function SymbolicIndexingInterface.is_param_sym(sys::AbstractSystem, sym)
+    !isnothing(SymbolicIndexingInterface.param_sym_to_index(sys, sym))
+end
+
 struct AbstractSysToExpr
     sys::AbstractSystem
     states::Vector

src/systems/jumps/jumpsystem.jl

@@ -354,7 +354,7 @@ function DiscreteProblemExpr(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing
         u0 = $u0
         p = $p
         tspan = $tspan
-        df = DiscreteFunction{true, true}(f, syms = $(Symbol.(states(sys))),
+        df = DiscreteFunction{true, true}(f; syms = $(Symbol.(states(sys))),
                                           indepsym = $(Symbol(get_iv(sys))),
                                           paramsyms = $(Symbol.(parameters(sys))))
         DiscreteProblem(df, u0, tspan, p)

test/odesystem.jl

@@ -930,9 +930,9 @@ let
     # TODO: maybe do not emit x_t
     sys4s = structural_simplify(sys4)
     prob = ODAEProblem(sys4s, [x => 1.0, D(x) => 1.0], (0, 1.0))
-    @test string.(prob.f.syms) == ["x(t)", "xˍt(t)"]
-    @test string.(prob.f.paramsyms) == ["pp"]
-    @test string(prob.f.indepsym) == "t"
+    @test string.(states(prob.f.sys)) == ["x(t)", "xˍt(t)"]
+    @test string.(parameters(prob.f.sys)) == ["pp"]
+    @test string.(independent_variables(prob.f.sys)) == ["t"]
 end
 
 let