Don't symbolic substitute in tearing

Author: YingboMa

src/bipartite_graph.jl

@@ -397,7 +397,6 @@ function Graphs.incidence_matrix(g::BipartiteGraph, val=true)
     S = sparse(I, J, val, nsrcs(g), ndsts(g))
 end
 
-
 """
     struct DiCMOBiGraph
 

src/structural_transformation/symbolics_tearing.jl

@@ -1,72 +1,114 @@
+"""
+    uneven_invmap(n::Int, list)
+
+returns an uneven inv map with length `n`.
+"""
+function uneven_invmap(n::Int, list)
+    rename = zero(Int, n)
+    for (i, v) in enumerate(list)
+        rename[v] = i
+    end
+    return rename
+end
+
+# N.B. assumes `slist` and `dlist` are unique
+function substitution_graph(graph, slist, dlist, var_eq_matching)
+    ns = length(slist)
+    nd = length(dlist)
+    ns == nd || error("internal error")
+    newgraph = BipartiteGraph(ns, nd)
+    erename = uneven_invmap(nsrc(graph), slist)
+    vrename = uneven_invmap(ndst(graph), dlist)
+    for e in 𝑠vertices(graph)
+        ie = erename[e]
+        ie == 0 && continue
+        for v in 𝑠neighbors(graph, e)
+            iv = vrename[v]
+            iv == 0 && continue
+            add_edge!(newgraph, ie, iv)
+        end
+    end
+
+    newmatching = zero(slist)
+    for (v, e) in enumerate(var_eq_matching)
+        iv = vrename[v]
+        ie = erename[e]
+        iv == 0 && continue
+        ie == 0 && error("internal error")
+        newmatching[iv] = ie
+    end
+
+    return newgraph, newmatching
+end
+
 function tearing_sub(expr, dict, s)
     expr = ModelingToolkit.fixpoint_sub(expr, dict)
     s ? simplify(expr) : expr
 end
 
+function tearing_substitution(sys::AbstractSystem; simplify=false)
+    (has_substitutions(sys) && !isnothing(get_substitutions(sys))) || return sys
+    subs = get_substitutions(sys)
+    neweqs = map(equations(sys)) do eq
+        if isdiffeq(eq)
+            return eq.lhs ~ tearing_sub(eq.rhs, solved, simplify)
+        else
+            if !(eq.lhs isa Number && eq.lhs == 0)
+                eq = 0 ~ eq.rhs - eq.lhs
+            end
+            rhs = tearing_sub(eq.rhs, solved, simplify)
+            if rhs isa Symbolic
+                return 0 ~ rhs
+            else # a number
+                error("tearing failled because the system is singular")
+            end
+        end
+        eq
+    end
+    @set! sys.eqs = neweqs
+end
+
+function solve_equation(eq, var, simplify)
+    rhs = value(solve_for(eq, var; simplify=simplify, check=false))
+    occursin(var, rhs) && error("solving $rhs for [$var] failed")
+    var ~ rhs
+end
+
+function normalize_equation(eq)
+    if !isdiffeq(eq)
+        if !(eq.lhs isa Number && eq.lhs == 0)
+            eq = 0 ~ eq.rhs - eq.lhs
+        end
+    end
+    eq
+end
+
 function tearing_reassemble(sys, var_eq_matching; simplify=false)
     s = structure(sys)
     @unpack fullvars, solvable_graph, graph = s
 
     eqs = equations(sys)
 
     ### extract partition information
-    function solve_equation(ieq, iv)
-        var = fullvars[iv]
-        eq = eqs[ieq]
-        rhs = value(solve_for(eq, var; simplify=simplify, check=false))
-
-        if var in vars(rhs)
-            # Usually we should be done here, but if we don't simplify we can get in
-            # trouble, so try our best to still solve for rhs
-            if !simplify
-                rhs = SymbolicUtils.polynormalize(rhs)
-            end
-
-            # Since we know `eq` is linear wrt `var`, so the round off must be a
-            # linear term. We can correct the round off error by a linear
-            # correction.
-            rhs -= expand_derivatives(Differential(var)(rhs))*var
-            (var in vars(rhs)) && throw(EquationSolveErrors(eq, var, rhs))
-        end
-        var => rhs
-    end
     is_solvable(eq, iv) = eq !== unassigned && BipartiteEdge(eq, iv) in solvable_graph
 
     solved_equations = Int[]
     solved_variables = Int[]
 
     # Solve solvable equations
-    for (iv, ieq) in enumerate(var_eq_matching);
-        is_solvable(ieq, iv) || continue
+    for (iv, ieq) in enumerate(var_eq_matching)
+        #is_solvable(ieq, iv) || continue
+        is_solvable(ieq, iv) || error("unreachable reached")
         push!(solved_equations, ieq); push!(solved_variables, iv)
     end
-
-    solved = Dict(solve_equation(ieq, iv) for (ieq, iv) in zip(solved_equations, solved_variables))
-    obseqs = [var ~ rhs for (var, rhs) in solved]
+    subgraph, submatching = substitution_graph(graph, slist, dlist, var_eq_matching)
+    toporder = topological_sort_by_dfs(DiCMOBiGraph{true}(subgraph, submatching))
+    substitutions = [solve_equation(eqs[solved_equations[i]], fullvars[solved_variables[i]], simplify) for i in toporder]
 
     # Rewrite remaining equations in terms of solved variables
-    function substitute_equation(ieq)
-        eq = eqs[ieq]
-        if isdiffeq(eq)
-            return eq.lhs ~ tearing_sub(eq.rhs, solved, simplify)
-        else
-            if !(eq.lhs isa Number && eq.lhs == 0)
-                eq = 0 ~ eq.rhs - eq.lhs
-            end
-            rhs = tearing_sub(eq.rhs, solved, simplify)
-            if rhs isa Symbolic
-                return 0 ~ rhs
-            else # a number
-                if abs(rhs) > 100eps(float(rhs))
-                    @warn "The equation $eq is not consistent. It simplifed to 0 == $rhs."
-                end
-                return nothing
-            end
-        end
-    end
 
-    neweqs = Any[substitute_equation(ieq) for ieq in 1:length(eqs) if !(ieq in solved_equations)]
-    filter!(!isnothing, neweqs)
+    solved_eq_set = BitSet(solved_equations)
+    neweqs = Equation[normalize_equation(eqs[ieq]) for ieq in 1:length(eqs) if !(ieq in solved_eq_set)]
 
     # Contract the vertices in the structure graph to make the structure match
     # the new reality of the system we've just created.
@@ -84,7 +126,7 @@ function tearing_reassemble(sys, var_eq_matching; simplify=false)
     @set! sys.structure = s
     @set! sys.eqs = neweqs
     @set! sys.states = [s.fullvars[idx] for idx in 1:length(s.fullvars) if !isdervar(s, idx)]
-    @set! sys.observed = [observed(sys); obseqs]
+    @set! sys.substitutions = substitutions
     return sys
 end
 

src/systems/abstractsystem.jl

@@ -222,6 +222,7 @@ for prop in [
              :connector_type
              :connections
              :preface
+             :substitutions
             ]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)

src/systems/diffeqs/odesystem.jl

@@ -96,16 +96,20 @@ struct ODESystem <: AbstractODESystem
     The integrator will use root finding to guarantee that it steps at each zero crossing.
     """
     continuous_events::Vector{SymbolicContinuousCallback}
+    """
+    substitutions: substitutions generated by tearing.
+    """
+    substitutions::Any
 
-    function ODESystem(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connector_type, connections, preface, events; checks::Bool = true)
+    function ODESystem(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connector_type, connections, preface, events, substitutions=nothing; checks::Bool=true)
         if checks
             check_variables(dvs,iv)
             check_parameters(ps,iv)
             check_equations(deqs,iv)
             check_equations(equations(events),iv)
             all_dimensionless([dvs;ps;iv]) || check_units(deqs)
         end
-        new(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connector_type, connections, preface, events)
+        new(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connector_type, connections, preface, events, substitutions)
     end
 end
 
@@ -153,7 +157,7 @@ function ODESystem(
         throw(ArgumentError("System names must be unique."))
     end
     cont_callbacks = SymbolicContinuousCallbacks(continuous_events)
-    ODESystem(deqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing, connector_type, nothing, preface, cont_callbacks, checks = checks)
+    ODESystem(deqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing, connector_type, nothing, preface, cont_callbacks, checks=checks)
 end
 
 function ODESystem(eqs, iv=nothing; kwargs...)

src/systems/diffeqs/sdesystem.jl

@@ -106,7 +106,7 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
                    defaults=_merge(Dict(default_u0), Dict(default_p)),
                    name=nothing,
                    connector_type=nothing,
-                   checks = true,
+                   checks=true,
                    )
     name === nothing && throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     deqs = scalarize(deqs)