Merge pull request #2249 from SciML/myb/syscompose

Fix `systems` keyword argument handling in `System`

Author: YingboMa

src/systems/abstractsystem.jl

@@ -472,15 +472,18 @@ function namespace_defaults(sys)
          for (k, v) in pairs(defs))
 end
 
-function namespace_equations(sys::AbstractSystem)
+function namespace_equations(sys::AbstractSystem, ivs = independent_variables(sys))
     eqs = equations(sys)
     isempty(eqs) && return Equation[]
-    map(eq -> namespace_equation(eq, sys), eqs)
+    map(eq -> namespace_equation(eq, sys; ivs), eqs)
 end
 
-function namespace_equation(eq::Equation, sys, n = nameof(sys))
-    _lhs = namespace_expr(eq.lhs, sys, n)
-    _rhs = namespace_expr(eq.rhs, sys, n)
+function namespace_equation(eq::Equation,
+    sys,
+    n = nameof(sys);
+    ivs = independent_variables(sys))
+    _lhs = namespace_expr(eq.lhs, sys, n; ivs)
+    _rhs = namespace_expr(eq.rhs, sys, n; ivs)
     _lhs ~ _rhs
 end
 
@@ -490,30 +493,29 @@ function namespace_assignment(eq::Assignment, sys)
     Assignment(_lhs, _rhs)
 end
 
-function namespace_expr(O, sys, n = nameof(sys))
-    ivs = independent_variables(sys)
+function namespace_expr(O, sys, n = nameof(sys); ivs = independent_variables(sys))
     O = unwrap(O)
     if any(isequal(O), ivs)
         return O
     elseif istree(O)
         T = typeof(O)
         renamed = let sys = sys, n = n, T = T
-            map(a -> namespace_expr(a, sys, n)::Any, arguments(O))
+            map(a -> namespace_expr(a, sys, n; ivs)::Any, arguments(O))
         end
         if isvariable(O)
             # Use renamespace so the scope is correct, and make sure to use the
             # metadata from the rescoped variable
             rescoped = renamespace(n, O)
             similarterm(O, operation(rescoped), renamed,
-                        metadata = metadata(rescoped))::T
+                metadata = metadata(rescoped))::T
         else
             similarterm(O, operation(O), renamed, metadata = metadata(O))::T
         end
     elseif isvariable(O)
         renamespace(n, O)
     elseif O isa Array
         let sys = sys, n = n
-            map(o -> namespace_expr(o, sys, n), O)
+            map(o -> namespace_expr(o, sys, n; ivs), O)
         end
     else
         O
@@ -661,20 +663,6 @@ function SymbolicIndexingInterface.is_param_sym(sys::AbstractSystem, sym)
     !isnothing(SymbolicIndexingInterface.param_sym_to_index(sys, sym))
 end
 
-struct AbstractSysToExpr
-    sys::AbstractSystem
-    states::Vector
-end
-AbstractSysToExpr(sys) = AbstractSysToExpr(sys, states(sys))
-function (f::AbstractSysToExpr)(O)
-    !istree(O) && return toexpr(O)
-    any(isequal(O), f.states) && return nameof(operation(O))  # variables
-    if issym(operation(O))
-        return build_expr(:call, Any[nameof(operation(O)); f.(arguments(O))])
-    end
-    return build_expr(:call, Any[operation(O); f.(arguments(O))])
-end
-
 ###
 ### System utils
 ###

src/systems/jumps/jumpsystem.jl

@@ -181,7 +181,6 @@ function generate_rate_function(js::JumpSystem, rate)
     end
     rf = build_function(rate, states(js), parameters(js),
         get_iv(js),
-        conv = states_to_sym(states(js)),
         expression = Val{true})
 end
 

src/systems/optimization/optimizationsystem.jl

@@ -124,7 +124,7 @@ function generate_gradient(sys::OptimizationSystem, vs = states(sys), ps = param
     grad = calculate_gradient(sys)
     pre = get_preprocess_constants(grad)
     return build_function(grad, vs, ps; postprocess_fbody = pre,
-        conv = AbstractSysToExpr(sys), kwargs...)
+        kwargs...)
 end
 
 function calculate_hessian(sys::OptimizationSystem)
@@ -140,14 +140,14 @@ function generate_hessian(sys::OptimizationSystem, vs = states(sys), ps = parame
     end
     pre = get_preprocess_constants(hess)
     return build_function(hess, vs, ps; postprocess_fbody = pre,
-        conv = AbstractSysToExpr(sys), kwargs...)
+        kwargs...)
 end
 
 function generate_function(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys);
     kwargs...)
     eqs = subs_constants(objective(sys))
     return build_function(eqs, vs, ps;
-        conv = AbstractSysToExpr(sys), kwargs...)
+        kwargs...)
 end
 
 function namespace_objective(sys::AbstractSystem)

src/systems/systems.jl

@@ -1,6 +1,6 @@
 function System(eqs::AbstractVector{<:Equation}, iv = nothing, args...; name = nothing,
     kw...)
-    ODESystem(eqs, iv, args...; name, checks = false)
+    ODESystem(eqs, iv, args...; name, kw..., checks = false)
 end
 
 """

src/utils.jl

@@ -50,30 +50,6 @@ end
 
 @deprecate substitute_expr!(expr, s) substitute(expr, s)
 
-function states_to_sym(states::Set)
-    function _states_to_sym(O)
-        if O isa Equation
-            Expr(:(=), _states_to_sym(O.lhs), _states_to_sym(O.rhs))
-        elseif istree(O)
-            op = operation(O)
-            args = arguments(O)
-            if issym(op)
-                O in states && return tosymbol(O)
-                # dependent variables
-                return build_expr(:call, Any[nameof(op); _states_to_sym.(args)])
-            else
-                canonical, O = canonicalexpr(O)
-                return canonical ? O : build_expr(:call, Any[op; _states_to_sym.(args)])
-            end
-        elseif O isa Num
-            return _states_to_sym(value(O))
-        else
-            return toexpr(O)
-        end
-    end
-end
-states_to_sym(states) = states_to_sym(Set(states))
-
 function todict(d)
     eltype(d) <: Pair || throw(ArgumentError("The variable-value mapping must be a Dict."))
     d isa Dict ? d : Dict(d)

test/dde.jl

@@ -83,26 +83,30 @@ sys = structural_simplify(sys)
 prob_mtk = SDDEProblem(sys, [x(t) => 1.0 + t], tspan; constant_lags = (τ,));
 @test_nowarn sol_mtk = solve(prob_mtk, RKMil())
 
-
 @variables t
 D = Differential(t)
 @parameters x(..) a
 
-function oscillator(;name, k=1.0, τ=0.01)
+function oscillator(; name, k = 1.0, τ = 0.01)
     @parameters k=k τ=τ
     @variables x(..)=0.1 y(t)=0.1 jcn(t)=0.0 delx(t)
     eqs = [D(x(t)) ~ y,
-           D(y) ~ -k*x(t-τ)+jcn,
-           delx ~ x(t-τ)]
-    return System(eqs; name=name)
+        D(y) ~ -k * x(t - τ) + jcn,
+        delx ~ x(t - τ)]
+    return System(eqs; name = name)
 end
 
-@named osc1 = oscillator(k=1.0, τ=0.01)
-@named osc2 = oscillator(k=2.0, τ=0.04)
+systems = @named begin
+    osc1 = oscillator(k = 1.0, τ = 0.01)
+    osc2 = oscillator(k = 2.0, τ = 0.04)
+end
 eqs = [osc1.jcn ~ osc2.delx,
-       osc2.jcn ~ osc1.delx]
+    osc2.jcn ~ osc1.delx]
 @named coupledOsc = System(eqs, t)
-@named coupledOsc = compose(coupledOsc, [osc1, osc2])
-sys = structural_simplify(coupledOsc)
-@test length(equations(sys)) == 4
-@test length(states(sys)) == 4
+@named coupledOsc = compose(coupledOsc, systems)
+@named coupledOsc2 = System(eqs, t; systems)
+for coupledOsc in [coupledOsc, coupledOsc2]
+    local sys = structural_simplify(coupledOsc)
+    @test length(equations(sys)) == 4
+    @test length(states(sys)) == 4
+end

test/optimizationsystem.jl

@@ -207,18 +207,20 @@ end
     prob = OptimizationProblem(combinedsys, u0, p, grad = true, hess = true, cons_j = true,
         cons_h = true)
     @test prob.f.sys === combinedsys
-    @test_broken SciMLBase.successful_retcode(solve(prob, Ipopt.Optimizer(); print_level = 0))
-   #=
-    @test sol.minimum < -1e5
-
-    prob = OptimizationProblem(sys2, [x => 0.0, y => 0.0], [a => 1.0, b => 100.0],
-        grad = true, hess = true, cons_j = true, cons_h = true)
-    @test prob.f.sys === sys2
-    sol = solve(prob, IPNewton())
-    @test sol.minimum < 1.0
-    sol = solve(prob, Ipopt.Optimizer(); print_level = 0)
-    @test sol.minimum < 1.0
-    =#
+    @test_broken SciMLBase.successful_retcode(solve(prob,
+        Ipopt.Optimizer();
+        print_level = 0))
+    #=
+     @test sol.minimum < -1e5
+
+     prob = OptimizationProblem(sys2, [x => 0.0, y => 0.0], [a => 1.0, b => 100.0],
+         grad = true, hess = true, cons_j = true, cons_h = true)
+     @test prob.f.sys === sys2
+     sol = solve(prob, IPNewton())
+     @test sol.minimum < 1.0
+     sol = solve(prob, Ipopt.Optimizer(); print_level = 0)
+     @test sol.minimum < 1.0
+     =#
 end
 
 @testset "metadata" begin