Merge pull request #1404 from SciML/myb/odaefix

Fix ODAEProblem codegen

Author: YingboMa

src/bipartite_graph.jl

@@ -2,7 +2,7 @@ module BipartiteGraphs
 
 export BipartiteEdge, BipartiteGraph, DiCMOBiGraph, Unassigned, unassigned,
         Matching, ResidualCMOGraph, InducedCondensationGraph, maximal_matching,
-        construct_augmenting_path!
+        construct_augmenting_path!, MatchedCondensationGraph
 
 export 𝑠vertices, 𝑑vertices, has_𝑠vertex, has_𝑑vertex, 𝑠neighbors, 𝑑neighbors,
        𝑠edges, 𝑑edges, nsrcs, ndsts, SRC, DST, set_neighbors!, invview,
@@ -516,4 +516,80 @@ end
 Graphs.has_edge(g::DiCMOBiGraph{true}, a, b) = a in inneighbors(g, b)
 Graphs.has_edge(g::DiCMOBiGraph{false}, a, b) = b in outneighbors(g, a)
 
+# Condensation Graphs
+abstract type AbstractCondensationGraph <: AbstractGraph{Int}; end
+function (T::Type{<:AbstractCondensationGraph})(g, sccs::Vector{Union{Int, Vector{Int}}})
+    scc_assignment = Vector{Int}(undef, isa(g, BipartiteGraph) ? ndsts(g) : nv(g))
+    for (i, c) in enumerate(sccs)
+        for v in c
+            scc_assignment[v] = i
+        end
+    end
+    T(g, sccs, scc_assignment)
+end
+(T::Type{<:AbstractCondensationGraph})(g, sccs::Vector{Vector{Int}}) =
+    T(g, Vector{Union{Int, Vector{Int}}}(sccs))
+
+Graphs.is_directed(::Type{<:AbstractCondensationGraph}) = true
+Graphs.nv(icg::AbstractCondensationGraph) = length(icg.sccs)
+Graphs.vertices(icg::AbstractCondensationGraph) = Base.OneTo(nv(icg))
+
+"""
+    struct MatchedCondensationGraph
+
+For some bipartite-graph and an orientation induced on its destination contraction,
+records the condensation DAG of the digraph formed by the orientation. I.e. this
+is a DAG of connected components formed by the destination vertices of some
+underlying bipartite graph.
+N.B.: This graph does not store explicit neighbor relations of the sccs.
+Therefor, the edge multiplicity is derived from the underlying bipartite graph,
+i.e. this graph is not strict.
+"""
+struct MatchedCondensationGraph{G <: DiCMOBiGraph} <: AbstractCondensationGraph
+    graph::G
+    # Records the members of a strongly connected component. For efficiency,
+    # trivial sccs (with one vertex member) are stored inline. Note: the sccs
+    # here need not be stored in topological order.
+    sccs::Vector{Union{Int, Vector{Int}}}
+    # Maps the vertices back to the scc of which they are a part
+    scc_assignment::Vector{Int}
+end
+
+
+Graphs.outneighbors(mcg::MatchedCondensationGraph, cc::Integer) =
+    Iterators.flatten((mcg.scc_assignment[v′] for v′ in outneighbors(mcg.graph, v) if mcg.scc_assignment[v′] != cc) for v in mcg.sccs[cc])
+
+Graphs.inneighbors(mcg::MatchedCondensationGraph, cc::Integer) =
+    Iterators.flatten((mcg.scc_assignment[v′] for v′ in inneighbors(mcg.graph, v) if mcg.scc_assignment[v′] != cc) for v in mcg.sccs[cc])
+
+"""
+    struct InducedCondensationGraph
+
+For some bipartite-graph and a topologicall sorted list of connected components,
+represents the condensation DAG of the digraph formed by the orientation. I.e. this
+is a DAG of connected components formed by the destination vertices of some
+underlying bipartite graph.
+N.B.: This graph does not store explicit neighbor relations of the sccs.
+Therefor, the edge multiplicity is derived from the underlying bipartite graph,
+i.e. this graph is not strict.
+"""
+struct InducedCondensationGraph{G <: BipartiteGraph} <: AbstractCondensationGraph
+    graph::G
+    # Records the members of a strongly connected component. For efficiency,
+    # trivial sccs (with one vertex member) are stored inline. Note: the sccs
+    # here are stored in topological order.
+    sccs::Vector{Union{Int, Vector{Int}}}
+    # Maps the vertices back to the scc of which they are a part
+    scc_assignment::Vector{Int}
+end
+
+_neighbors(icg::InducedCondensationGraph, cc::Integer) =
+    Iterators.flatten(Iterators.flatten(icg.graph.fadjlist[vsrc] for vsrc in icg.graph.badjlist[v]) for v in icg.sccs[cc])
+
+Graphs.outneighbors(icg::InducedCondensationGraph, v::Integer) =
+    (icg.scc_assignment[n] for n in _neighbors(icg, v) if icg.scc_assignment[n] > v)
+
+Graphs.inneighbors(icg::InducedCondensationGraph, v::Integer) =
+    (icg.scc_assignment[n] for n in _neighbors(icg, v) if icg.scc_assignment[n] < v)
+
 end # module

src/compat/incremental_cycles.jl

@@ -1,4 +1,5 @@
 using Base.Iterators: repeated
+import Graphs.Experimental.Traversals: topological_sort
 
 # Abstract Interface
 

src/structural_transformation/codegen.jl

@@ -4,7 +4,7 @@ using ModelingToolkit: isdifferenceeq, has_continuous_events, generate_rootfindi
 
 const MAX_INLINE_NLSOLVE_SIZE = 8
 
-function torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs)
+function torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs, nlsolve_scc_idxs, states_idxs)
     s = structure(sys)
     @unpack fullvars, graph = s
 
@@ -42,54 +42,52 @@ function torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs)
     # from previous partitions. Hence, we can build the dependency chain as we
     # traverse the partitions.
 
-    # `avars2dvars` maps a algebraic variable to its differential variable
-    # dependencies.
-    avars2dvars = Dict{Int,Set{Int}}()
-    c = 0
-    for scc in var_sccs
-        v_residual = scc
-        e_residual = [var_eq_matching[c] for c in v_residual if var_eq_matching[c] !== unassigned]
-        # initialization
-        for tvar in v_residual
-            avars2dvars[tvar] = Set{Int}()
+    var_rename = ones(Int64, ndsts(graph))
+    nlsolve_vars = Int[]
+    for i in nlsolve_scc_idxs, c in var_sccs[i]
+        append!(nlsolve_vars, c)
+        for v in c
+            var_rename[v] = 0
         end
-        for teq in e_residual
-            c += 1
-            for var in 𝑠neighbors(graph, teq)
-                # Skip the tearing variables in the current partition, because
-                # we are computing them from all the other states.
-                Graphs.insorted(var, v_residual) && continue
-                deps = get(avars2dvars, var, nothing)
-                if deps === nothing # differential variable
-                    @assert !isalgvar(s, var)
-                    for tvar in v_residual
-                        push!(avars2dvars[tvar], var)
-                    end
-                else # tearing variable from previous partitions
-                    @assert isalgvar(s, var)
-                    for tvar in v_residual
-                        union!(avars2dvars[tvar], avars2dvars[var])
-                    end
-                end
+    end
+    masked_cumsum!(var_rename)
+
+    dig = DiCMOBiGraph{true}(graph, var_eq_matching)
+
+    fused_var_deps = map(1:ndsts(graph)) do v
+        BitSet(var_rename[v′] for v′ in neighborhood(dig, v, Inf; dir=:in) if var_rename[v′] != 0)
+    end
+
+    for scc in var_sccs
+        if length(scc) >= 2
+            deps = fused_var_deps[scc[1]]
+            for c in 2:length(scc)
+                union!(deps, fused_var_deps[c])
             end
         end
     end
 
-    dvrange = diffvars_range(s)
-    dvar2idx = Dict(v=>i for (i, v) in enumerate(dvrange))
+    nlsolve_eqs = BitSet(var_eq_matching[c]::Int for c in nlsolve_vars if var_eq_matching[c] !== unassigned)
+
+    var2idx = Dict(v => i for (i, v) in enumerate(states_idxs))
+    nlsolve_vars_set = BitSet(nlsolve_vars)
+
     I = Int[]; J = Int[]
     eqidx = 0
     for ieq in 𝑠vertices(graph)
-        isalgeq(s, ieq) && continue
+        ieq in nlsolve_eqs && continue
         eqidx += 1
         for ivar in 𝑠neighbors(graph, ieq)
-            if isdiffvar(s, ivar)
+            isdervar(s, ivar) && continue
+            if var_rename[ivar] != 0
                 push!(I, eqidx)
-                push!(J, dvar2idx[ivar])
-            elseif isalgvar(s, ivar)
-                for dvar in avars2dvars[ivar]
+                push!(J, var2idx[ivar])
+            else
+                for dvar in fused_var_deps[ivar]
+                    isdervar(s, dvar) && continue
+                    dvar in nlsolve_vars_set && continue
                     push!(I, eqidx)
-                    push!(J, dvar2idx[dvar])
+                    push!(J, var2idx[dvar])
                 end
             end
         end
@@ -123,7 +121,17 @@ function gen_nlsolve(eqs, vars, u0map::AbstractDict; checkbounds=true)
     params = setdiff(allvars, vars) # these are not the subject of the root finding
 
     # splatting to tighten the type
-    u0 = [map(var->get(u0map, var, 1e-3), vars)...]
+    u0 = []
+    for v in vars
+        v in keys(u0map) || (push!(u0, 1e-3); continue)
+        u = substitute(v, u0map)
+        for i in 1:length(u0map)
+            u = substitute(u, u0map)
+            u isa Number && (push!(u0, u); break)
+        end
+        u isa Number || error("$v doesn't have a default.")
+    end
+    u0 = [u0...]
     # specialize on the scalar case
     isscalar = length(u0) == 1
     u0 = isscalar ? u0[1] : SVector(u0...)
@@ -175,23 +183,30 @@ function build_torn_function(
     s = structure(sys)
     @unpack fullvars = s
     var_eq_matching, var_sccs = algebraic_variables_scc(sys)
+    condensed_graph = MatchedCondensationGraph(
+        DiCMOBiGraph{true}(complete(s.graph), complete(var_eq_matching)), var_sccs)
+    toporder = topological_sort_by_dfs(condensed_graph)
+    var_sccs = var_sccs[toporder]
 
-    states = map(i->s.fullvars[i], diffvars_range(s))
-    mass_matrix_diag = ones(length(states))
+    states_idxs = collect(diffvars_range(s))
+    mass_matrix_diag = ones(length(states_idxs))
     torn_expr = []
     defs = defaults(sys)
+    nlsolve_scc_idxs = Int[]
 
     needs_extending = false
-    for scc in var_sccs
-        torn_vars = [s.fullvars[var] for var in scc if var_eq_matching[var] !== unassigned]
-        torn_eqs = [eqs[var_eq_matching[var]] for var in scc if var_eq_matching[var] !== unassigned]
+    for (i, scc) in enumerate(var_sccs)
+        #torn_vars = [s.fullvars[var] for var in scc if var_eq_matching[var] !== unassigned]
+        torn_vars_idxs = Int[var for var in scc if var_eq_matching[var] !== unassigned]
+        torn_eqs = [eqs[var_eq_matching[var]] for var in torn_vars_idxs]
         isempty(torn_eqs) && continue
         if length(torn_eqs) <= max_inlining_size
-            append!(torn_expr, gen_nlsolve(torn_eqs, torn_vars, defs, checkbounds=checkbounds))
+            append!(torn_expr, gen_nlsolve(torn_eqs, s.fullvars[torn_vars_idxs], defs, checkbounds=checkbounds))
+            push!(nlsolve_scc_idxs, i)
         else
             needs_extending = true
             append!(rhss, map(x->x.rhs, torn_eqs))
-            append!(states, torn_vars)
+            append!(states_idxs, torn_vars_idxs)
             append!(mass_matrix_diag, zeros(length(torn_eqs)))
         end
     end
@@ -205,7 +220,8 @@ function build_torn_function(
         rhss
     )
 
-    syms = map(Symbol, states)
+    states = s.fullvars[states_idxs]
+    syms = map(Symbol, states_idxs)
     pre = get_postprocess_fbody(sys)
 
     expr = SymbolicUtils.Code.toexpr(
@@ -237,7 +253,7 @@ function build_torn_function(
 
         ODEFunction{true}(
                           @RuntimeGeneratedFunction(expr),
-                          sparsity = torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs),
+                          sparsity = jacobian_sparsity ? torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs, nlsolve_scc_idxs, states_idxs) : nothing,
                           syms = syms,
                           observed = observedfun,
                           mass_matrix = mass_matrix,

src/structural_transformation/tearing.jl

@@ -38,19 +38,24 @@ function contract_variables(graph::BipartiteGraph, var_eq_matching::Matching, el
         [var_rename[v′] for v′ in neighborhood(dig, v, Inf; dir=:in) if var_rename[v′] != 0]
     end
 
-    new_fadjlist = Vector{Int}[
-        let new_list = Vector{Int}()
-            for v in graph.fadjlist[i]
-                if var_rename[v] != 0
-                    push!(new_list, var_rename[v])
-                else
-                    append!(new_list, var_deps[v])
+    nelim = length(eliminated_variables)
+    newgraph = BipartiteGraph(nsrcs(graph) - nelim, ndsts(graph) - nelim)
+    for e in 𝑠vertices(graph)
+        ne = eq_rename[e]
+        ne == 0 && continue
+        for v in 𝑠neighbors(graph, e)
+            newvar = var_rename[v]
+            if newvar != 0
+                add_edge!(newgraph, ne, newvar)
+            else
+                for nv in var_deps[v]
+                    add_edge!(newgraph, ne, nv)
                 end
             end
-            new_list
-        end for i = 1:nsrcs(graph) if eq_rename[i] != 0]
+        end
+    end
 
-    return BipartiteGraph(new_fadjlist, ndsts(graph) - length(eliminated_variables))
+    return newgraph
 end
 
 """