Merge pull request #2175 from SciML/myb/optsys

Fix `structural_transformation` for optimization systems

Author: YingboMa

src/structural_transformation/bipartite_tearing/modia_tearing.jl

@@ -73,7 +73,10 @@ function tear_graph_modia(structure::SystemStructure, isder::F = nothing,
     # find them here [TODO: It would be good to have an explicit example of this.]
 
     @unpack graph, solvable_graph = structure
-    var_eq_matching = complete(maximal_matching(graph, eqfilter, varfilter, U))
+    var_eq_matching = maximal_matching(graph, eqfilter, varfilter, U)
+    var_eq_matching = complete(var_eq_matching,
+                               max(length(var_eq_matching),
+                                   maximum(x -> x isa Int ? x : 0, var_eq_matching)))
     var_sccs::Vector{Union{Vector{Int}, Int}} = find_var_sccs(graph, var_eq_matching)
     vargraph = DiCMOBiGraph{true}(graph)
     ict = IncrementalCycleTracker(vargraph; dir = :in)

src/structural_transformation/symbolics_tearing.jl

@@ -583,10 +583,10 @@ Tear the nonlinear equations in system. When `simplify=true`, we simplify the
 new residual equations after tearing. End users are encouraged to call [`structural_simplify`](@ref)
 instead, which calls this function internally.
 """
-function tearing(sys::AbstractSystem; simplify = false)
-    state = TearingState(sys)
+function tearing(sys::AbstractSystem, state = TearingState(sys); mm = nothing,
+                 simplify = false, kwargs...)
     var_eq_matching = tearing(state)
-    invalidate_cache!(tearing_reassemble(state, var_eq_matching; simplify = simplify))
+    invalidate_cache!(tearing_reassemble(state, var_eq_matching; mm, simplify))
 end
 
 """

src/systems/alias_elimination.jl

@@ -468,10 +468,10 @@ function observed2graph(eqs, states)
 end
 
 function fixpoint_sub(x, dict)
-    y = substitute(x, dict)
+    y = fast_substitute(x, dict)
     while !isequal(x, y)
         y = x
-        x = substitute(y, dict)
+        x = fast_substitute(y, dict)
     end
 
     return x

src/systems/optimization/optimizationsystem.jl

@@ -609,18 +609,19 @@ function structural_simplify(sys::OptimizationSystem; kwargs...)
         end
     end
     nlsys = NonlinearSystem(econs, states(sys), parameters(sys); name = :___tmp_nlsystem)
-    snlsys = structural_simplify(nlsys; check_consistency = false, kwargs...)
+    snlsys = structural_simplify(nlsys; fully_determined = false, kwargs...)
     obs = observed(snlsys)
     subs = Dict(eq.lhs => eq.rhs for eq in observed(snlsys))
     seqs = equations(snlsys)
-    cons_simplified = Array{eltype(cons), 1}(undef, length(icons) + length(seqs))
+    cons_simplified = similar(cons, length(icons) + length(seqs))
     for (i, eq) in enumerate(Iterators.flatten((seqs, icons)))
-        cons_simplified[i] = substitute(eq, subs)
+        cons_simplified[i] = fixpoint_sub(eq, subs)
     end
     newsts = setdiff(states(sys), keys(subs))
     @set! sys.constraints = cons_simplified
     @set! sys.observed = [observed(sys); obs]
-    @set! sys.op = substitute(equations(sys), subs)
+    neweqs = fixpoint_sub.(equations(sys), (subs,))
+    @set! sys.op = length(neweqs) == 1 ? first(neweqs) : neweqs
     @set! sys.states = newsts
     return sys
 end

src/systems/systemstructure.jl

@@ -554,14 +554,15 @@ function merge_io(io, inputs)
 end
 
 function structural_simplify!(state::TearingState, io = nothing; simplify = false,
-                              check_consistency = true, kwargs...)
+                              check_consistency = true, fully_determined = true,
+                              kwargs...)
     if state.sys isa ODESystem
         ci = ModelingToolkit.ClockInference(state)
         ModelingToolkit.infer_clocks!(ci)
         tss, inputs, continuous_id, id_to_clock = ModelingToolkit.split_system(ci)
         cont_io = merge_io(io, inputs[continuous_id])
         sys, input_idxs = _structural_simplify!(tss[continuous_id], cont_io; simplify,
-                                                check_consistency,
+                                                check_consistency, fully_determined,
                                                 kwargs...)
         if length(tss) > 1
             # TODO: rename it to something else
@@ -576,7 +577,7 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
                 end
                 dist_io = merge_io(io, inputs[i])
                 ss, = _structural_simplify!(state, dist_io; simplify, check_consistency,
-                                            kwargs...)
+                                            fully_determined, kwargs...)
                 append!(appended_parameters, inputs[i], states(ss))
                 discrete_subsystems[i] = ss
             end
@@ -588,14 +589,16 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
         end
     else
         sys, input_idxs = _structural_simplify!(state, io; simplify, check_consistency,
-                                                kwargs...)
+                                                fully_determined, kwargs...)
     end
     has_io = io !== nothing
     return has_io ? (sys, input_idxs) : sys
 end
 
 function _structural_simplify!(state::TearingState, io; simplify = false,
-                               check_consistency = true, kwargs...)
+                               check_consistency = true, fully_determined = true,
+                               kwargs...)
+    check_consistency &= fully_determined
     has_io = io !== nothing
     orig_inputs = Set()
     if has_io
@@ -606,7 +609,11 @@ function _structural_simplify!(state::TearingState, io; simplify = false,
     if check_consistency
         ModelingToolkit.check_consistency(state, orig_inputs)
     end
-    sys = ModelingToolkit.dummy_derivative(sys, state; simplify, mm)
+    if fully_determined
+        sys = ModelingToolkit.dummy_derivative(sys, state; simplify, mm, check_consistency)
+    else
+        sys = ModelingToolkit.tearing(sys, state; simplify, mm, check_consistency)
+    end
     fullstates = [map(eq -> eq.lhs, observed(sys)); states(sys)]
     @set! sys.observed = ModelingToolkit.topsort_equations(observed(sys), fullstates)
     ModelingToolkit.invalidate_cache!(sys), input_idxs

src/utils.jl

@@ -845,13 +845,19 @@ end
 
 # Symbolics needs to call unwrap on the substitution rules, but most of the time
 # we don't want to do that in MTK.
-function fast_substitute(eq::Equation, subs)
-    fast_substitute(eq.lhs, subs) ~ fast_substitute(eq.rhs, subs)
+const Eq = Union{Equation, Inequality}
+function fast_substitute(eq::Eq, subs)
+    if eq isa Inequality
+        Inequality(fast_substitute(eq.lhs, subs), fast_substitute(eq.rhs, subs),
+                   eq.relational_op)
+    else
+        Equation(fast_substitute(eq.lhs, subs), fast_substitute(eq.rhs, subs))
+    end
 end
-function fast_substitute(eq::Equation, subs::Pair)
-    fast_substitute(eq.lhs, subs) ~ fast_substitute(eq.rhs, subs)
+function fast_substitute(eq::T, subs::Pair) where {T <: Eq}
+    T(fast_substitute(eq.lhs, subs), fast_substitute(eq.rhs, subs))
 end
-fast_substitute(eqs::AbstractArray{Equation}, subs) = fast_substitute.(eqs, (subs,))
+fast_substitute(eqs::AbstractArray{<:Eq}, subs) = fast_substitute.(eqs, (subs,))
 fast_substitute(a, b) = substitute(a, b)
 function fast_substitute(expr, pair::Pair)
     a, b = pair