Many fixes

Author: YingboMa

src/ModelingToolkit.jl

@@ -42,7 +42,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,7 +152,7 @@ for prop in [
              :iv
              :states
              :ps
-             :array_vars
+             :var_to_name
              :ctrls
              :defaults
              :observed
@@ -202,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
@@ -234,29 +238,17 @@ 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
 
-    avs = get_array_vars(sys)
+    avs = get_var_to_name(sys)
     v = get(avs, name, nothing)
-    v === nothing || return namespace ? renamespace(sysname, v) : v
+    v === nothing || return namespace ? renamespace(sysname, v, name) : v
 
     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_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
-
     if has_observed(sys)
         obs = get_observed(sys)
         i = findfirst(x->getname(x.lhs)==name,obs)
@@ -301,13 +293,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

@@ -32,7 +32,7 @@ struct ODESystem <: AbstractODESystem
     """Parameter variables. Must not contain the independent variable."""
     ps::Vector
     """Array variables."""
-    array_vars
+    var_to_name
     """Control parameters (some subset of `ps`)."""
     ctrls::Vector
     """Observed states."""
@@ -84,11 +84,11 @@ struct ODESystem <: AbstractODESystem
     """
     connection_type::Any
 
-    function ODESystem(deqs, iv, dvs, ps, array_vars, 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, array_vars, 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
 
@@ -121,9 +121,9 @@ function ODESystem(
     dvs′ = value.(scalarize(dvs))
     ps′ = value.(scalarize(ps))
 
-    array_vars = Dict()
-    process_variables!(array_vars, defaults, dvs′)
-    process_variables!(array_vars, 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))
@@ -134,7 +134,7 @@ function ODESystem(
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    ODESystem(deqs, iv′, dvs′, ps′, array_vars, 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
 
 vars(x::Sym) = Set([x])

src/systems/diffeqs/sdesystem.jl

@@ -38,7 +38,7 @@ struct SDESystem <: AbstractODESystem
     """Parameter variables. Must not contain the independent variable."""
     ps::Vector
     """Array variables."""
-    array_vars
+    var_to_name
     """Control parameters (some subset of `ps`)."""
     ctrls::Vector
     """Observed states."""
@@ -86,11 +86,11 @@ struct SDESystem <: AbstractODESystem
     """
     connection_type::Any
 
-    function SDESystem(deqs, neqs, iv, dvs, ps, array_vars, 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, array_vars, 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
 
@@ -120,16 +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))
 
-    array_vars = Dict()
-    process_variables!(array_vars, defaults, dvs′)
-    process_variables!(array_vars, 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′, array_vars, 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

@@ -31,7 +31,7 @@ struct DiscreteSystem <: AbstractSystem
     """Parameter variables. Must not contain the independent variable."""
     ps::Vector
     """Array variables."""
-    array_vars
+    var_to_name
     """Control parameters (some subset of `ps`)."""
     ctrls::Vector
     """Observed states."""
@@ -54,10 +54,10 @@ struct DiscreteSystem <: AbstractSystem
     in `DiscreteSystem`.
     """
     default_p::Dict
-    function DiscreteSystem(discreteEqs, iv, dvs, ps, array_vars, 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, array_vars, 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
 
@@ -88,15 +88,15 @@ function DiscreteSystem(
     defaults = todict(defaults)
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
 
-    array_vars = Dict()
-    process_variables!(array_vars, defaults, dvs′)
-    process_variables!(array_vars, 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′, array_vars, 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

@@ -38,7 +38,7 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractSystem
     """The parameters of the system. Must not contain the independent variable."""
     ps::Vector
     """Array variables."""
-    array_vars
+    var_to_name
     observed::Vector{Equation}
     """The name of the system. . These are required to have unique names."""
     name::Symbol
@@ -53,10 +53,10 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractSystem
     type: type of the system
     """
     connection_type::Any
-    function JumpSystem{U}(ap::U, iv, states, ps, array_vars, 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, array_vars, observed, name, systems, defaults, connection_type)
+        new{U}(ap, iv, states, ps, var_to_name, observed, name, systems, defaults, connection_type)
     end
 end
 
@@ -94,11 +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)
-    array_vars = Dict()
-    process_variables!(array_vars, defaults, dvs′)
-    process_variables!(array_vars, defaults, ps′)
+    var_to_name = Dict()
+    process_variables!(var_to_name, defaults, dvs′)
+    process_variables!(var_to_name, defaults, ps′)
 
-    JumpSystem{typeof(ap)}(ap, value(iv), states, ps, array_vars, 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

@@ -26,7 +26,7 @@ struct NonlinearSystem <: AbstractSystem
     """Parameters."""
     ps::Vector
     """Array variables."""
-    array_vars
+    var_to_name
     observed::Vector{Equation}
     """
     Jacobian matrix. Note: this field will not be defined until
@@ -78,11 +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)
-    array_vars = Dict()
-    process_variables!(array_vars, defaults, dvs′)
-    process_variables!(array_vars, defaults, ps′)
+    var_to_name = Dict()
+    process_variables!(var_to_name, defaults, states)
+    process_variables!(var_to_name, defaults, ps)
 
-    NonlinearSystem(eqs, states, ps, array_vars, 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

@@ -24,7 +24,7 @@ struct OptimizationSystem <: AbstractSystem
     """Parameters."""
     ps::Vector
     """Array variables."""
-    array_vars
+    var_to_name
     observed::Vector{Equation}
     equality_constraints::Vector{Equation}
     inequality_constraints::Vector
@@ -63,12 +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)
-    array_vars = Dict()
-    process_variables!(array_vars, defaults, dvs′)
-    process_variables!(array_vars, 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, array_vars,
+                       value(op), states, ps, var_to_name,
                        observed,
                        equality_constraints, inequality_constraints,
                        name, systems, defaults

src/utils.jl

@@ -160,9 +160,9 @@ 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!(array_vars, defs, vars)
+function process_variables!(var_to_name, defs, vars)
     collect_defaults!(defs, vars)
-    collect_array_vars!(array_vars, vars)
+    collect_var_to_name!(var_to_name, vars)
     return nothing
 end
 
@@ -173,26 +173,26 @@ function collect_defaults!(defs, vars)
     return defs
 end
 
-function array_vars(x, vars=Dict{Symbol, Any}())
+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->array_vars(a, v), arguments(x))
-            array_vars(operation(x), v)
+            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
-                array_vars(x.expr, vars)
+                var_to_name(x.expr, vars)
             end
         else
-            array_vars(operation(x), vars)
+            var_to_name(operation(x), vars)
             for a in arguments(x)
-                array_vars(a, vars)
+                var_to_name(a, vars)
             end
         end
     elseif x isa Sym && symtype(x) <: AbstractArray
@@ -202,9 +202,9 @@ function array_vars(x, vars=Dict{Symbol, Any}())
     vars
 end
 
-function collect_array_vars!(vars, xs)
+function collect_var_to_name!(vars, xs)
     for x in xs
-        ax = array_vars(x)
+        ax = var_to_name(x)
         if isempty(ax)
             vars[getname(x)] = x
         else