Merge pull request #1125 from SciML/myb/arrayvar

`getproperty` for array variables

Author: YingboMa

src/ModelingToolkit.jl

@@ -43,7 +43,7 @@ using Reexport
 @reexport using Symbolics
 export @derivatives
 using Symbolics: _parse_vars, value, @derivatives, get_variables,
-                 exprs_occur_in, solve_for, build_expr
+                 exprs_occur_in, solve_for, build_expr, unwrap, wrap
 import Symbolics: rename, get_variables!, _solve, hessian_sparsity,
                   jacobian_sparsity, islinear, _iszero, _isone,
                   tosymbol, lower_varname, diff2term, var_from_nested_derivative,

src/systems/abstractsystem.jl

@@ -152,6 +152,7 @@ for prop in [
              :iv
              :states
              :ps
+             :var_to_name
              :ctrls
              :defaults
              :observed
@@ -201,6 +202,10 @@ Setfield.get(obj::AbstractSystem, ::Setfield.PropertyLens{field}) where {field}
 end
 
 rename(x::AbstractSystem, name) = @set x.name = name
+function rename(xx::Symbolics.ArrayOp, name)
+    @set! xx.expr.f.arguments[1] = rename(xx.expr.f.arguments[1], name)
+    @set! xx.term.arguments[2] = rename(xx.term.arguments[2], name)
+end
 
 function Base.propertynames(sys::AbstractSystem; private=false)
     if private
@@ -233,25 +238,34 @@ function getvar(sys::AbstractSystem, name::Symbol; namespace=false)
     elseif !isempty(systems)
         i = findfirst(x->nameof(x)==name, systems)
         if i !== nothing
-            return namespace ? rename(systems[i], renamespace(sysname,name)) : systems[i]
+            return namespace ? rename(systems[i], renamespace(sysname, name)) : systems[i]
         end
     end
 
-    sts = get_states(sys)
-    i = findfirst(x->getname(x) == name, sts)
-
-    if i !== nothing
-        return namespace ? renamespace(sysname,sts[i]) : sts[i]
-    end
+    if has_var_to_name(sys)
+        avs = get_var_to_name(sys)
+        v = get(avs, name, nothing)
+        v === nothing || return namespace ? renamespace(sysname, v, name) : v
 
-    if has_ps(sys)
-        ps = get_ps(sys)
-        i = findfirst(x->getname(x) == name,ps)
+    else
+        sts = get_states(sys)
+        i = findfirst(x->getname(x) == name, sts)
         if i !== nothing
-            return namespace ? renamespace(sysname,ps[i]) : ps[i]
+            return namespace ? renamespace(sysname,sts[i]) : sts[i]
+        end
+
+        if has_ps(sys)
+            ps = get_ps(sys)
+            i = findfirst(x->getname(x) == name,ps)
+            if i !== nothing
+                return namespace ? renamespace(sysname,ps[i]) : ps[i]
+            end
         end
     end
 
+    sts = get_states(sys)
+    i = findfirst(x->getname(x) == name, sts)
+
     if has_observed(sys)
         obs = get_observed(sys)
         i = findfirst(x->getname(x.lhs)==name,obs)
@@ -296,13 +310,12 @@ ParentScope(sym::Union{Num, Symbolic}) = setmetadata(sym, SymScope, ParentScope(
 struct GlobalScope <: SymScope end
 GlobalScope(sym::Union{Num, Symbolic}) = setmetadata(sym, SymScope, GlobalScope())
 
-function renamespace(namespace, x)
-    if x isa Num
-        renamespace(namespace, value(x))
-    elseif x isa Symbolic
+function renamespace(namespace, x, name=nothing)
+    x = unwrap(x)
+    if x isa Symbolic
         let scope = getmetadata(x, SymScope, LocalScope())
             if scope isa LocalScope
-                rename(x, renamespace(namespace, getname(x)))
+                rename(x, renamespace(namespace, name === nothing ? getname(x) : name))
             elseif scope isa ParentScope
                 setmetadata(x, SymScope, scope.parent)
             else # GlobalScope

src/systems/diffeqs/odesystem.jl

@@ -31,6 +31,8 @@ struct ODESystem <: AbstractODESystem
     states::Vector
     """Parameter variables. Must not contain the independent variable."""
     ps::Vector
+    """Array variables."""
+    var_to_name
     """Control parameters (some subset of `ps`)."""
     ctrls::Vector
     """Observed states."""
@@ -82,11 +84,11 @@ struct ODESystem <: AbstractODESystem
     """
     connection_type::Any
 
-    function ODESystem(deqs, iv, dvs, ps, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type)
+    function ODESystem(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type)
         check_variables(dvs,iv)
         check_parameters(ps,iv)
         check_equations(deqs,iv)
-        new(deqs, iv, dvs, ps, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type)
+        new(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type)
     end
 end
 
@@ -119,8 +121,9 @@ function ODESystem(
     dvs′ = value.(scalarize(dvs))
     ps′ = value.(scalarize(ps))
 
-    collect_defaults!(defaults, dvs′)
-    collect_defaults!(defaults, ps′)
+    var_to_name = Dict()
+    process_variables!(var_to_name, defaults, dvs′)
+    process_variables!(var_to_name, defaults, ps′)
 
     tgrad = RefValue(Vector{Num}(undef, 0))
     jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
@@ -131,7 +134,7 @@ function ODESystem(
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    ODESystem(deqs, iv′, dvs′, ps′, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing, connection_type)
+    ODESystem(deqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing, connection_type)
 end
 
 function ODESystem(eqs, iv=nothing; kwargs...)

src/systems/diffeqs/sdesystem.jl

@@ -37,6 +37,8 @@ struct SDESystem <: AbstractODESystem
     states::Vector
     """Parameter variables. Must not contain the independent variable."""
     ps::Vector
+    """Array variables."""
+    var_to_name
     """Control parameters (some subset of `ps`)."""
     ctrls::Vector
     """Observed states."""
@@ -84,11 +86,11 @@ struct SDESystem <: AbstractODESystem
     """
     connection_type::Any
 
-    function SDESystem(deqs, neqs, iv, dvs, ps, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
+    function SDESystem(deqs, neqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
         check_variables(dvs,iv)
         check_parameters(ps,iv)
         check_equations(deqs,iv)
-        new(deqs, neqs, iv, dvs, ps, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
+        new(deqs, neqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
     end
 end
 
@@ -118,15 +120,16 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
     defaults = todict(defaults)
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
 
-    collect_defaults!(defaults, dvs′)
-    collect_defaults!(defaults, ps′)
+    var_to_name = Dict()
+    process_variables!(var_to_name, defaults, dvs′)
+    process_variables!(var_to_name, defaults, ps′)
 
     tgrad = RefValue(Vector{Num}(undef, 0))
     jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
     ctrl_jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
     Wfact   = RefValue(Matrix{Num}(undef, 0, 0))
     Wfact_t = RefValue(Matrix{Num}(undef, 0, 0))
-    SDESystem(deqs, neqs, iv′, dvs′, ps′, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
+    SDESystem(deqs, neqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
 end
 
 function generate_diffusion_function(sys::SDESystem, dvs = states(sys), ps = parameters(sys); kwargs...)

src/systems/discrete_system/discrete_system.jl

@@ -30,6 +30,8 @@ struct DiscreteSystem <: AbstractSystem
     states::Vector
     """Parameter variables. Must not contain the independent variable."""
     ps::Vector
+    """Array variables."""
+    var_to_name
     """Control parameters (some subset of `ps`)."""
     ctrls::Vector
     """Observed states."""
@@ -52,10 +54,10 @@ struct DiscreteSystem <: AbstractSystem
     in `DiscreteSystem`.
     """
     default_p::Dict
-    function DiscreteSystem(discreteEqs, iv, dvs, ps, ctrls, observed, name, systems, default_u0, default_p)
+    function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, default_u0, default_p)
         check_variables(dvs,iv)
         check_parameters(ps,iv)
-        new(discreteEqs, iv, dvs, ps, ctrls, observed, name, systems, default_u0, default_p)
+        new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, default_u0, default_p)
     end
 end
 
@@ -86,14 +88,15 @@ function DiscreteSystem(
     defaults = todict(defaults)
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
 
-    collect_defaults!(defaults, dvs′)
-    collect_defaults!(defaults, ps′)
+    var_to_name = Dict()
+    process_variables!(var_to_name, defaults, dvs′)
+    process_variables!(var_to_name, defaults, ps′)
 
     sysnames = nameof.(systems)
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    DiscreteSystem(eqs, iv′, dvs′, ps′, ctrl′, observed, name, systems, default_u0, default_p)
+    DiscreteSystem(eqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, name, systems, default_u0, default_p)
 end
 
 """

src/systems/jumps/jumpsystem.jl

@@ -37,6 +37,8 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractSystem
     states::Vector
     """The parameters of the system. Must not contain the independent variable."""
     ps::Vector
+    """Array variables."""
+    var_to_name
     observed::Vector{Equation}
     """The name of the system. . These are required to have unique names."""
     name::Symbol
@@ -51,10 +53,10 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractSystem
     type: type of the system
     """
     connection_type::Any
-    function JumpSystem{U}(ap::U, iv, states, ps, observed, name, systems, defaults, connection_type) where U <: ArrayPartition
+    function JumpSystem{U}(ap::U, iv, states, ps, var_to_name, observed, name, systems, defaults, connection_type) where U <: ArrayPartition
         check_variables(states, iv)
         check_parameters(ps, iv)
-        new{U}(ap, iv, states, ps, observed, name, systems, defaults, connection_type)
+        new{U}(ap, iv, states, ps, var_to_name, observed, name, systems, defaults, connection_type)
     end
 end
 
@@ -92,10 +94,11 @@ function JumpSystem(eqs, iv, states, ps;
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
 
     states, ps = value.(states), value.(ps)
-    collect_defaults!(defaults, states)
-    collect_defaults!(defaults, ps)
+    var_to_name = Dict()
+    process_variables!(var_to_name, defaults, states)
+    process_variables!(var_to_name, defaults, ps)
 
-    JumpSystem{typeof(ap)}(ap, value(iv), states, ps, observed, name, systems, defaults, connection_type)
+    JumpSystem{typeof(ap)}(ap, value(iv), states, ps, var_to_name, observed, name, systems, defaults, connection_type)
 end
 
 function generate_rate_function(js, rate)

src/systems/nonlinear/nonlinearsystem.jl

@@ -25,6 +25,8 @@ struct NonlinearSystem <: AbstractSystem
     states::Vector
     """Parameters."""
     ps::Vector
+    """Array variables."""
+    var_to_name
     observed::Vector{Equation}
     """
     Jacobian matrix. Note: this field will not be defined until
@@ -76,10 +78,11 @@ function NonlinearSystem(eqs, states, ps;
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
 
     states, ps = value.(states), value.(ps)
-    collect_defaults!(defaults, states)
-    collect_defaults!(defaults, ps)
+    var_to_name = Dict()
+    process_variables!(var_to_name, defaults, states)
+    process_variables!(var_to_name, defaults, ps)
 
-    NonlinearSystem(eqs, states, ps, observed, jac, name, systems, defaults, nothing, connection_type)
+    NonlinearSystem(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, nothing, connection_type)
 end
 
 function calculate_jacobian(sys::NonlinearSystem; sparse=false, simplify=false)

src/systems/optimization/optimizationsystem.jl

@@ -23,6 +23,8 @@ struct OptimizationSystem <: AbstractSystem
     states::Vector
     """Parameters."""
     ps::Vector
+    """Array variables."""
+    var_to_name
     observed::Vector{Equation}
     equality_constraints::Vector{Equation}
     inequality_constraints::Vector
@@ -61,11 +63,12 @@ function OptimizationSystem(op, states, ps;
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
 
     states, ps = value.(states), value.(ps)
-    collect_defaults!(defaults, states)
-    collect_defaults!(defaults, ps)
+    var_to_name = Dict()
+    process_variables!(var_to_name, defaults, states)
+    process_variables!(var_to_name, defaults, ps)
 
     OptimizationSystem(
-                       value(op), states, ps,
+                       value(op), states, ps, var_to_name,
                        observed,
                        equality_constraints, inequality_constraints,
                        name, systems, defaults

src/utils.jl

@@ -160,19 +160,66 @@ hasdefault(v) = hasmetadata(v, Symbolics.VariableDefaultValue)
 getdefault(v) = value(getmetadata(v, Symbolics.VariableDefaultValue))
 setdefault(v, val) = val === nothing ? v : setmetadata(v, Symbolics.VariableDefaultValue, value(val))
 
+function process_variables!(var_to_name, defs, vars)
+    collect_defaults!(defs, vars)
+    collect_var_to_name!(var_to_name, vars)
+    return nothing
+end
+
 function collect_defaults!(defs, vars)
     for v in vars; (haskey(defs, v) || !hasdefault(v)) && continue
         defs[v] = getdefault(v)
     end
     return defs
 end
 
+function var_to_name(x, vars=Dict{Symbol, Any}())
+    x = Symbolics.unwrap(x)
+    if istree(x)
+        if hasmetadata(x, Symbolics.GetindexParent)
+            v = Dict{Symbol, Any}()
+            foreach(a->var_to_name(a, v), arguments(x))
+            var_to_name(operation(x), v)
+            name = first(only(v))
+            vars[name] = getmetadata(x, Symbolics.GetindexParent)
+        elseif x isa Symbolics.ArrayOp
+            t = x.term
+            if istree(t) && operation(t) === (map) && arguments(t)[1] isa Symbolics.CallWith
+                vars[nameof(arguments(t)[2])] = x
+            else
+                var_to_name(x.expr, vars)
+            end
+        else
+            var_to_name(operation(x), vars)
+            for a in arguments(x)
+                var_to_name(a, vars)
+            end
+        end
+    elseif x isa Sym && symtype(x) <: AbstractArray
+        vars[nameof(x)] = x
+    end
+
+    vars
+end
+
+function collect_var_to_name!(vars, xs)
+    for x in xs
+        ax = var_to_name(x)
+        if isempty(ax)
+            vars[getname(x)] = x
+        else
+            merge!(vars, ax)
+        end
+    end
+    return vars
+end
+
 "Throw error when difference/derivative operation occurs in the R.H.S."
 @noinline function throw_invalid_operator(opvar, eq, op::Type)
     if op === Difference
         optext = "difference"
     elseif op === Differential
-        optext="derivative" 
+        optext="derivative"
     end
     msg = "The $optext variable must be isolated to the left-hand " *
     "side of the equation like `$opvar ~ ...`.\n Got $eq."

test/controlsystem.jl

@@ -28,4 +28,4 @@ sol = solve(prob,BFGS())
     sys1 = ControlSystem(loss,eqs_short, t, [x, v], [u], p, name = :sys1)
     sys2 = ControlSystem(loss,eqs_short, t, [x, v], [u], p, name = :sys1)
     @test_throws ArgumentError ControlSystem(loss, [sys2.v ~ sys1.v], t, [], [], [], systems = [sys1, sys2])
-end
+end