Merge pull request #1157 from SciML/sm/refactor

More robust array symbolic handling

Author: YingboMa

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "5.26.0"
+version = "6.0.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -73,7 +73,7 @@ Setfield = "0.7"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
 SymbolicUtils = "0.12, 0.13"
-Symbolics = "1.4.1"
+Symbolics = "2.0"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"
 julia = "1.2"

examples/rc_model.jl

@@ -14,4 +14,5 @@ rc_eqs = [
           connect(capacitor.n, source.n, ground.g)
          ]
 
-@named rc_model = compose(ODESystem(rc_eqs, t), [resistor, capacitor, source, ground])
+@named rc_model = ODESystem(rc_eqs, t)
+rc_model = compose(rc_model, [resistor, capacitor, source, ground])

examples/serial_inductor.jl

@@ -13,4 +13,5 @@ eqs = [
        connect(source.n, inductor2.n, ground.g)
       ]
 
-@named ll_model = compose(ODESystem(eqs, t), source, resistor, inductor1, inductor2, ground)
+@named ll_model = ODESystem(eqs, t)
+ll_model = compose(ll_model, [source, resistor, inductor1, inductor2, ground])

src/ModelingToolkit.jl

@@ -33,7 +33,7 @@ using RecursiveArrayTools
 import SymbolicUtils
 import SymbolicUtils: istree, arguments, operation, similarterm, promote_symtype,
                       Symbolic, Term, Add, Mul, Pow, Sym, FnType,
-                      @rule, Rewriters, substitute
+                      @rule, Rewriters, substitute, metadata
 using SymbolicUtils.Code
 import SymbolicUtils.Code: toexpr
 import SymbolicUtils.Rewriters: Chain, Postwalk, Prewalk, Fixpoint
@@ -43,7 +43,8 @@ using Reexport
 @reexport using Symbolics
 export @derivatives
 using Symbolics: _parse_vars, value, @derivatives, get_variables,
-                 exprs_occur_in, solve_for, build_expr, unwrap, wrap
+                 exprs_occur_in, solve_for, build_expr, unwrap, wrap,
+                 VariableSource, getname, variable
 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

@@ -131,13 +131,6 @@ function generate_function end
 
 Base.nameof(sys::AbstractSystem) = getfield(sys, :name)
 
-function getname(t)
-    if istree(t)
-        operation(t) isa Sym ? getname(operation(t)) : error("Cannot get name of $t")
-    else
-        nameof(t)
-    end
-end
 #Deprecated
 function independent_variable(sys::AbstractSystem)
     Base.depwarn("`independent_variable` is deprecated. Use `get_iv` or `independent_variables` instead.",:independent_variable)
@@ -193,6 +186,7 @@ for prop in [
              :depvars
              :indvars
              :connection_type
+             :preface
             ]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)
@@ -248,37 +242,35 @@ function Base.propertynames(sys::AbstractSystem; private=false)
     end
 end
 
-Base.getproperty(sys::AbstractSystem, name::Symbol; namespace=true) = getvar(sys, name; namespace=namespace)
+Base.getproperty(sys::AbstractSystem, name::Symbol; namespace=true) = wrap(getvar(sys, name; namespace=namespace))
 function getvar(sys::AbstractSystem, name::Symbol; namespace=false)
-    sysname = nameof(sys)
     systems = get_systems(sys)
     if isdefined(sys, name)
         Base.depwarn("`sys.name` like `sys.$name` is deprecated. Use getters like `get_$name` instead.", "sys.$name")
         return getfield(sys, name)
     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(sys, name)) : systems[i]
         end
     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
-
+        v === nothing || return namespace ? renamespace(sys, v) : v
     else
         sts = get_states(sys)
         i = findfirst(x->getname(x) == name, sts)
         if i !== nothing
-            return namespace ? renamespace(sysname,sts[i]) : sts[i]
+            return namespace ? renamespace(sys, 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]
+                return namespace ? renamespace(sys, ps[i]) : ps[i]
             end
         end
     end
@@ -290,7 +282,7 @@ function getvar(sys::AbstractSystem, name::Symbol; namespace=false)
         obs = get_observed(sys)
         i = findfirst(x->getname(x.lhs)==name,obs)
         if i !== nothing
-            return namespace ? renamespace(sysname,obs[i]) : obs[i]
+            return namespace ? renamespace(sys, obs[i]) : obs[i]
         end
     end
 
@@ -330,20 +322,39 @@ 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, name=nothing)
+renamespace(sys, eq::Equation) = namespace_equation(eq, sys)
+
+function _renamespace(sys, x)
+    v = unwrap(x)
+
+    if istree(v) && symtype(operation(v)) <: FnType
+        ov = metadata(operation(v), metadata(v))
+        return similarterm(v, renamespace(sys, ov), arguments(v), symtype(v), metadata=metadata(v))
+    end
+
+    if v isa Namespace
+        sysp, v = v.parent, v.named
+        sysn = Symbol(getname(sys), :., getname(sysp))
+        sys = sys isa AbstractSystem ? rename(sysp, sysn) : sysn
+    end
+
+    Namespace(sys, v)
+end
+
+function renamespace(sys, x)
     x = unwrap(x)
     if x isa Symbolic
         let scope = getmetadata(x, SymScope, LocalScope())
             if scope isa LocalScope
-                rename(x, renamespace(namespace, name === nothing ? getname(x) : name))
+                _renamespace(sys, x)
             elseif scope isa ParentScope
                 setmetadata(x, SymScope, scope.parent)
             else # GlobalScope
                 x
             end
         end
     else
-        Symbol(namespace,:₊,x)
+        Symbol(getname(sys), :., x)
     end
 end
 
@@ -353,33 +364,38 @@ namespace_controls(sys::AbstractSystem) = controls(sys, controls(sys))
 
 function namespace_defaults(sys)
     defs = defaults(sys)
-    Dict((isparameter(k) ? parameters(sys, k) : states(sys, k)) => namespace_expr(defs[k], nameof(sys), independent_variables(sys)) for k in keys(defs))
+    Dict((isparameter(k) ? parameters(sys, k) : states(sys, k)) => namespace_expr(defs[k], sys) for k in keys(defs))
 end
 
 function namespace_equations(sys::AbstractSystem)
     eqs = equations(sys)
     isempty(eqs) && return Equation[]
-    ivs = independent_variables(sys)
-    map(eq -> namespace_equation(eq, nameof(sys), ivs), eqs)
+    map(eq->namespace_equation(eq, sys), eqs)
 end
 
-function namespace_equation(eq::Equation, name, ivs)
-    _lhs = namespace_expr(eq.lhs, name, ivs)
-    _rhs = namespace_expr(eq.rhs, name, ivs)
+function namespace_equation(eq::Equation, sys)
+    _lhs = namespace_expr(eq.lhs, sys)
+    _rhs = namespace_expr(eq.rhs, sys)
     _lhs ~ _rhs
 end
 
-function namespace_expr(O::Sym, name, ivs)
-    any(isequal(O), ivs) ? O : renamespace(name, O)
+function namespace_assignment(eq::Assignment, sys)
+    _lhs = namespace_expr(eq.lhs, sys)
+    _rhs = namespace_expr(eq.rhs, sys)
+    Assignment(_lhs, _rhs)
 end
 
-_symparam(s::Symbolic{T}) where {T} = T
-function namespace_expr(O, name, ivs) where {T}
-    O = value(O)
-    if istree(O)
-        renamed = map(a -> namespace_expr(a, name, ivs), arguments(O))
-        if operation(O) isa Sym
-            renamespace(name, O)
+function namespace_expr(O, sys) where {T}
+    ivs = independent_variables(sys)
+    O = unwrap(O)
+    if any(isequal(O), ivs)
+        return O
+    elseif isvariable(O)
+        renamespace(sys, O)
+    elseif istree(O)
+        renamed = map(a->namespace_expr(a, sys), arguments(O))
+        if symtype(operation(O)) <: FnType
+            renamespace(sys, O)
         else
             similarterm(O, operation(O), renamed)
         end
@@ -409,13 +425,12 @@ function controls(sys::AbstractSystem)
 end
 
 function observed(sys::AbstractSystem)
-    ivs = independent_variables(sys)
     obs = get_observed(sys)
     systems = get_systems(sys)
     [obs;
      reduce(vcat,
-            (map(o -> namespace_equation.(o, nameof(s), ivs), observed(s)) for s in systems),
-            init = Equation[])]
+            (map(o->namespace_equation(o, s), observed(s)) for s in systems),
+            init=Equation[])]
 end
 
 Base.@deprecate default_u0(x) defaults(x) false
@@ -426,7 +441,7 @@ function defaults(sys::AbstractSystem)
     isempty(systems) ? defs : mapreduce(namespace_defaults, merge, systems; init=defs)
 end
 
-states(sys::AbstractSystem, v) = renamespace(nameof(sys), v)
+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)
@@ -448,6 +463,25 @@ function equations(sys::ModelingToolkit.AbstractSystem)
     end
 end
 
+function preface(sys::ModelingToolkit.AbstractSystem)
+    has_preface(sys) || return nothing
+    pre = get_preface(sys)
+    systems = get_systems(sys)
+    if isempty(systems)
+        return pre
+    else
+        pres = pre === nothing ? [] : pre
+        for sys in systems
+            pre = get_preface(sys)
+            pre === nothing && continue
+            for eq in pre
+                push!(pres, namespace_assignment(eq, sys))
+            end
+        end
+        return isempty(pres) ? nothing : pres
+    end
+end
+
 function islinear(sys::AbstractSystem)
     rhs = [eq.rhs for eq ∈ equations(sys)]
 
@@ -913,7 +947,11 @@ function Base.hash(sys::AbstractSystem, s::UInt)
     s = foldr(hash, get_systems(sys), init=s)
     s = foldr(hash, get_states(sys), init=s)
     s = foldr(hash, get_ps(sys), init=s)
-    s = foldr(hash, get_eqs(sys), init=s)
+    if sys isa OptimizationSystem
+        s = hash(get_op(sys), s)
+    else
+        s = foldr(hash, get_eqs(sys), init=s)
+    end
     s = foldr(hash, get_observed(sys), init=s)
     s = hash(independent_variables(sys), s)
     return s
@@ -935,7 +973,7 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem; name::Symbol=nameo
     else
         throw("Extending multivariate systems is not supported")
     end
-    
+
     eqs = union(equations(basesys), equations(sys))
     sts = union(states(basesys), states(sys))
     ps = union(parameters(basesys), parameters(sys))

src/systems/control/controlsystem.jl

@@ -87,7 +87,8 @@ function ControlSystem(loss, deqs::AbstractVector{<:Equation}, iv, dvs, controls
                        default_u0=Dict(),
                        default_p=Dict(),
                        defaults=_merge(Dict(default_u0), Dict(default_p)),
-                       name=gensym(:ControlSystem))
+                       name=nothing)
+    name === nothing && throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     if !(isempty(default_u0) && isempty(default_p))
         Base.depwarn("`default_u0` and `default_p` are deprecated. Use `defaults` instead.", :ControlSystem, force=true)
     end
@@ -165,7 +166,7 @@ function runge_kutta_discretize(sys::ControlSystem,dt,tspan;
     L = @RuntimeGeneratedFunction(build_function(lo,sts,ctr,ps,iv,conv = ModelingToolkit.ControlToExpr(sys)))
 
     var(n, i...) = var(nameof(n), i...)
-    var(n::Symbol, i...) = Sym{FnType{Tuple{symtype(iv)}, Real}}(nameof(Variable(n, i...)))
+    var(n::Symbol, i...) = variable(n, i..., T=FnType)
     # Expand out all of the variables in time and by stages
     timed_vars = [[var(operation(x),i)(iv) for i in 1:n+1] for x in states(sys)]
     k_vars = [[var(Symbol(:ᵏ,nameof(operation(x))),i,j)(iv) for i in 1:m, j in 1:n] for x in states(sys)]
@@ -196,5 +197,5 @@ function runge_kutta_discretize(sys::ControlSystem,dt,tspan;
     equalities = vcat(stages,updates,control_equality)
     opt_states = vcat(reduce(vcat,reduce(vcat,states_timeseries)),reduce(vcat,reduce(vcat,k_timeseries)),reduce(vcat,reduce(vcat,control_timeseries)))
 
-    OptimizationSystem(reduce(+,losses, init=0),opt_states,ps,equality_constraints = equalities)
+    OptimizationSystem(reduce(+,losses, init=0),opt_states,ps,equality_constraints = equalities, name=nameof(sys))
 end

src/systems/diffeqs/abstractodesystem.jl

@@ -10,7 +10,7 @@ function calculate_tgrad(sys::AbstractODESystem;
   xs = states(sys)
   rule = Dict(map((x, xt) -> xt=>x, detime_dvs.(xs), xs))
   rhs = substitute.(rhs, Ref(rule))
-  tgrad = [expand_derivatives(ModelingToolkit.Differential(iv)(r), simplify) for r in rhs]
+  tgrad = [expand_derivatives(Differential(iv)(r), simplify) for r in rhs]
   reverse_rule = Dict(map((x, xt) -> x=>xt, detime_dvs.(xs), xs))
   tgrad = Num.(substitute.(tgrad, Ref(reverse_rule)))
   get_tgrad(sys)[] = tgrad
@@ -102,10 +102,16 @@ function generate_function(
     p = map(x->time_varying_as_func(value(x), sys), ps)
     t = get_iv(sys)
 
+    if has_preface(sys) && (pre = preface(sys); pre !== nothing)
+      pre = ex -> Let(pre, ex)
+    else
+        pre = ex -> ex
+    end
+
     if implicit_dae
-        build_function(rhss, ddvs, u, p, t; kwargs...)
+        build_function(rhss, ddvs, u, p, t; postprocess_fbody=pre, kwargs...)
     else
-        build_function(rhss, u, p, t; kwargs...)
+        build_function(rhss, u, p, t; postprocess_fbody=pre, kwargs...)
     end
 end
 
@@ -409,13 +415,13 @@ function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
     =#
 
     fsym = gensym(:f)
-    _f = :($fsym = ModelingToolkit.ODEFunctionClosure($f_oop, $f_iip))
+    _f = :($fsym = $ODEFunctionClosure($f_oop, $f_iip))
     tgradsym = gensym(:tgrad)
     if tgrad
         tgrad_oop, tgrad_iip = generate_tgrad(sys, dvs, ps;
                                 simplify=simplify,
                                 expression=Val{true}, kwargs...)
-        _tgrad = :($tgradsym = ModelingToolkit.ODEFunctionClosure($tgrad_oop, $tgrad_iip))
+        _tgrad = :($tgradsym = $ODEFunctionClosure($tgrad_oop, $tgrad_iip))
     else
         _tgrad = :($tgradsym = nothing)
     end
@@ -425,7 +431,7 @@ function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
         jac_oop,jac_iip = generate_jacobian(sys, dvs, ps;
                                  sparse=sparse, simplify=simplify,
                                  expression=Val{true}, kwargs...)
-        _jac = :($jacsym = ModelingToolkit.ODEFunctionClosure($jac_oop, $jac_iip))
+        _jac = :($jacsym = $ODEFunctionClosure($jac_oop, $jac_iip))
     else
         _jac = :($jacsym = nothing)
     end