Fix state selection

Author: Keno

src/bipartite_graph.jl

@@ -306,6 +306,25 @@ function Graphs.add_edge!(g::BipartiteGraph, edge::BipartiteEdge, md=NO_METADATA
     return true  # edge successfully added
 end
 
+Graphs.rem_edge!(g::BipartiteGraph, i::Integer, j::Integer) =
+    Graphs.rem_edge!(g, BipartiteEdge(i, j))
+function Graphs.rem_edge!(g::BipartiteGraph, edge::BipartiteEdge)
+    @unpack fadjlist, badjlist = g
+    s, d = src(edge), dst(edge)
+    (has_𝑠vertex(g, s) && has_𝑑vertex(g, d)) || error("edge ($edge) out of range.")
+    @inbounds list = fadjlist[s]
+    index = searchsortedfirst(list, d)
+    @inbounds (index <= length(list) && list[index] == d) || error("graph does not have edge $edge")
+    deleteat!(list, index)
+    g.ne -= 1
+    if badjlist isa AbstractVector
+        @inbounds list = badjlist[d]
+        index = searchsortedfirst(list, s)
+        deleteat!(list, index)
+    end
+    return true  # edge successfully deleted
+end
+
 function Graphs.add_vertex!(g::BipartiteGraph{T}, type::VertType) where T
     if type === DST
         if g.badjlist isa AbstractVector

src/structural_transformation/bipartite_tearing/modia_tearing.jl

@@ -13,6 +13,14 @@
 #
 ################################################
 
+function try_assign_eq!(ict::IncrementalCycleTracker, vj::Integer, eq::Integer)
+    G = ict.graph
+    add_edge_checked!(ict, Iterators.filter(!=(vj), 𝑠neighbors(G.graph, eq)), vj) do G
+        G.matching[vj] = eq
+        G.ne += length(𝑠neighbors(G.graph, eq)) - 1
+    end
+end
+
 """
     (eSolved, vSolved, eResidue, vTear) = tearEquations!(td, Gsolvable, es, vs)
 
@@ -33,10 +41,7 @@ function tearEquations!(ict::IncrementalCycleTracker, Gsolvable, es::Vector{Int}
     for eq in es  # iterate only over equations that are not in eSolvedFixed
         for vj in Gsolvable[eq]
             if G.matching[vj] === unassigned && (vj in vActive)
-                r = add_edge_checked!(ict, Iterators.filter(!=(vj), 𝑠neighbors(G.graph, eq)), vj) do G
-                    G.matching[vj] = eq
-                    G.ne += length(𝑠neighbors(G.graph, eq)) - 1
-                end
+                r = try_assign_eq!(ict, vj, eq)
                 r && break
             end
         end
@@ -45,6 +50,15 @@ function tearEquations!(ict::IncrementalCycleTracker, Gsolvable, es::Vector{Int}
     return ict
 end
 
+function tear_graph_block_modia!(var_eq_matching, graph, solvable_graph, eqs, vars)
+    ict = IncrementalCycleTracker(DiCMOBiGraph{true}(graph); dir=:in)
+    tearEquations!(ict, solvable_graph.fadjlist, eqs, vars)
+    for var in vars
+        var_eq_matching[var] = ict.graph.matching[var]
+    end
+    return nothing
+end
+
 """
     tear_graph_modia(sys) -> sys
 
@@ -58,13 +72,10 @@ function tear_graph_modia(graph::BipartiteGraph, solvable_graph::BipartiteGraph;
     for vars in var_sccs
         filtered_vars = filter(varfilter, vars)
         ieqs = Int[var_eq_matching[v] for v in filtered_vars if var_eq_matching[v] !== unassigned]
-
-        ict = IncrementalCycleTracker(DiCMOBiGraph{true}(graph); dir=:in)
-        tearEquations!(ict, solvable_graph.fadjlist, ieqs, filtered_vars)
-
         for var in vars
-            var_eq_matching[var] = ict.graph.matching[var]
+            var_eq_matching[var] = unassigned
         end
+        tear_graph_block_modia!(var_eq_matching, graph, solvable_graph, ieqs, filtered_vars)
     end
 
     return var_eq_matching

src/structural_transformation/pantelides.jl

@@ -75,17 +75,33 @@ function pantelides!(sys::ODESystem; maxiters = 8000)
     find_solvables!(sys)
     s = structure(sys)
     # D(j) = assoc[j]
-    @unpack graph, var_to_diff, solvable_graph = s
-    return (sys, pantelides!(graph, solvable_graph, var_to_diff)...)
+    @unpack graph, var_to_diff = s
+    # N.B.: var_derivative! and eq_derivative! are defined in symbolics_tearing.jl
+    return (sys, pantelides!(PantelidesSetup(sys, graph, var_to_diff))...)
 end
 
-function pantelides!(graph, solvable_graph, var_to_diff; maxiters = 8000)
+struct PantelidesSetup{T}
+    system::T
+    graph::BipartiteGraph
+    var_to_diff::DiffGraph
+    eq_to_diff::DiffGraph
+    var_eq_matching::Matching
+end
+
+function PantelidesSetup(sys::T, graph, var_to_diff) where {T}
     neqs = nsrcs(graph)
     nvars = nv(var_to_diff)
-    vcolor = falses(nvars)
-    ecolor = falses(neqs)
     var_eq_matching = Matching(nvars)
     eq_to_diff = DiffGraph(neqs)
+    PantelidesSetup{T}(sys, graph, var_to_diff, eq_to_diff, var_eq_matching)
+end
+
+function pantelides!(p::PantelidesSetup; maxiters = 8000)
+    @unpack graph, var_to_diff, eq_to_diff, var_eq_matching = p
+    neqs = nsrcs(graph)
+    nvars = nv(var_to_diff)
+    vcolor = falses(nvars)
+    ecolor = falses(neqs)
     neqs′ = neqs
     for k in 1:neqs′
         eq′ = k
@@ -107,27 +123,22 @@ function pantelides!(graph, solvable_graph, var_to_diff; maxiters = 8000)
             for var in eachindex(vcolor); vcolor[var] || continue
                 # introduce a new variable
                 nvars += 1
-                add_vertex!(graph, DST); add_vertex!(solvable_graph, DST)
+                add_vertex!(graph, DST);
                 # the new variable is the derivative of `var`
 
                 add_edge!(var_to_diff, var, add_vertex!(var_to_diff))
                 push!(var_eq_matching, unassigned)
+                var_derivative!(p, eq)
             end
 
             for eq in eachindex(ecolor); ecolor[eq] || continue
                 # introduce a new equation
                 neqs += 1
-                add_vertex!(graph, SRC); add_vertex!(solvable_graph, SRC)
+                add_vertex!(graph, SRC);
                 # the new equation is created by differentiating `eq`
                 eq_diff = add_vertex!(eq_to_diff)
                 add_edge!(eq_to_diff, eq, eq_diff)
-                for var in 𝑠neighbors(graph, eq)
-                    add_edge!(graph, eq_diff, var)
-                    add_edge!(graph, eq_diff, var_to_diff[var])
-                    # If you have f(x) = 0, then the derivative is (∂f/∂x) ẋ = 0.
-                    # which is linear, thus solvable in ẋ.
-                    add_edge!(solvable_graph, eq_diff, var_to_diff[var])
-                end
+                eq_derivative!(p, eq)
             end
 
             for var in eachindex(vcolor); vcolor[var] || continue

src/structural_transformation/partial_state_selection.jl

@@ -2,21 +2,108 @@ function partial_state_selection_graph!(sys::ODESystem)
     s = get_structure(sys)
     (s isa SystemStructure) || (sys = initialize_system_structure(sys))
     s = structure(sys)
-    find_solvables!(sys)
+    find_solvables!(sys; allow_symbolic=true)
     @set! s.graph = complete(s.graph)
     @set! sys.structure = s
-    (sys, partial_state_selection_graph!(s.graph, s.solvable_graph, s.var_to_diff)...)
+    var_eq_matching, eq_to_diff = pantelides!(PantelidesSetup(sys, s.graph, s.var_to_diff))
+    (sys, partial_state_selection_graph!(s.graph, s.solvable_graph, s.var_to_diff, var_eq_matching, eq_to_diff)...)
+end
+
+function ascend_dg(xs, dg, level)
+    while level > 0
+        xs = Int[dg[x] for x in xs]
+        level -= 1
+    end
+    return xs
+end
+
+function ascend_dg_all(xs, dg, level, maxlevel)
+    r = Int[]
+    while true
+        if level <= 0
+            append!(r, xs)
+        end
+        maxlevel <= 0 && break
+        xs = Int[dg[x] for x in xs if dg[x] !== nothing]
+        level -= 1
+        maxlevel -= 1
+    end
+    return r
+end
+
+function pss_graph_modia!(graph, solvable_graph, var_eq_matching, var_to_diff, eq_to_diff, varlevel, inv_varlevel, inv_eqlevel)
+    # var_eq_matching is a maximal matching on the top-differentiated variables.
+    # Find Strongly connected components. Note that after pantelides, we expect
+    # a balanced system, so a maximal matching should be possible.
+    var_sccs::Vector{Union{Vector{Int}, Int}} = find_var_sccs(graph, var_eq_matching)
+    var_eq_matching = Matching{Union{Unassigned, SelectedState}}(var_eq_matching)
+    for vars in var_sccs
+        # TODO: We should have a way to not have the scc code look at unassigned vars.
+        if length(vars) == 1 && varlevel[vars[1]] != 0
+            continue
+        end
+
+        # Now proceed level by level from lowest to highest and tear the graph.
+        eqs = [var_eq_matching[var] for var in vars]
+        maxlevel = level = maximum(map(x->inv_eqlevel[x], eqs))
+        old_level_vars = ()
+        ict = IncrementalCycleTracker(DiCMOBiGraph{true}(graph, complete(Matching(ndsts(graph)))); dir=:in)
+        while level >= 0
+            to_tear_eqs_toplevel = filter(eq->inv_eqlevel[eq] >= level, eqs)
+            to_tear_eqs = ascend_dg(to_tear_eqs_toplevel, invview(eq_to_diff), level)
+
+            to_tear_vars_toplevel = filter(var->inv_varlevel[var] >= level, vars)
+            to_tear_vars = ascend_dg_all(to_tear_vars_toplevel, invview(var_to_diff), level, maxlevel)
+
+            if old_level_vars !== ()
+                # Inherit constraints from previous level.
+                # TODO: Is this actually a good idea or do we want full freedom
+                # to tear differently on each level? Does it make a difference
+                # whether we're using heuristic or optimal tearing?
+                removed_eqs = Int[]
+                removed_vars = Int[]
+                for var in old_level_vars
+                    old_assign = ict.graph.matching[var]
+                    if !isa(old_assign, Int) || ict.graph.matching[var_to_diff[var]] !== unassigned
+                        continue
+                    end
+                    # Make sure the ict knows about this edge, so it doesn't accidentally introduce
+                    # a cycle.
+                    ok = try_assign_eq!(ict, var_to_diff[var], eq_to_diff[old_assign])
+                    @assert ok
+                    var_eq_matching[var_to_diff[var]] = eq_to_diff[old_assign]
+                    push!(removed_eqs, eq_to_diff[ict.graph.matching[var]])
+                    push!(removed_vars, var_to_diff[var])
+                end
+                to_tear_eqs = setdiff(to_tear_eqs, removed_eqs)
+                to_tear_vars = setdiff(to_tear_vars, removed_vars)
+            end
+            filter!(var->ict.graph.matching[var] === unassigned, to_tear_vars)
+            filter!(eq->invview(ict.graph.matching)[eq] === unassigned, to_tear_eqs)
+            tearEquations!(ict, solvable_graph.fadjlist, to_tear_eqs, to_tear_vars)
+            for var in to_tear_vars
+                var_eq_matching[var] = ict.graph.matching[var]
+            end
+            old_level_vars = to_tear_vars
+            level -= 1
+        end
+        for var in old_level_vars
+            if varlevel[var] !== 0 && var_eq_matching[var] === unassigned
+                var_eq_matching[var] = SelectedState()
+            end
+        end
+    end
+    return var_eq_matching
 end
 
 struct SelectedState; end
-function partial_state_selection_graph!(graph, solvable_graph, var_to_diff)
-    var_eq_matching, eq_to_diff = pantelides!(graph, solvable_graph, var_to_diff)
+function partial_state_selection_graph!(graph, solvable_graph, var_to_diff, var_eq_matching, eq_to_diff)
     eq_to_diff = complete(eq_to_diff)
 
-    eqlevel = map(1:nsrcs(graph)) do eq
+    inv_eqlevel = map(1:nsrcs(graph)) do eq
         level = 0
-        while eq_to_diff[eq] !== nothing
-            eq = eq_to_diff[eq]
+        while invview(eq_to_diff)[eq] !== nothing
+            eq = invview(eq_to_diff)[eq]
             level += 1
         end
         level
@@ -31,45 +118,25 @@ function partial_state_selection_graph!(graph, solvable_graph, var_to_diff)
         level
     end
 
-    all_selected_states = Int[]
-
-    level = 0
-    level_vars = [var for var in 1:ndsts(graph) if varlevel[var] == 0 && invview(var_to_diff)[var] !== nothing]
+    inv_varlevel = map(1:ndsts(graph)) do var
+        level = 0
+        while invview(var_to_diff)[var] !== nothing
+            var = invview(var_to_diff)[var]
+            level += 1
+        end
+        level
+    end
 
-    # TODO: Is this actually useful or should we just compute another maximal matching?
+    # TODO: Should pantelides just return this?
     for var in 1:ndsts(graph)
-        if !(var in level_vars)
+        if var_to_diff[var] !== nothing
             var_eq_matching[var] = unassigned
         end
     end
 
-    while level < maximum(eqlevel)
-        var_eq_matching = tear_graph_modia(graph, solvable_graph;
-            eqfilter  = eq->eqlevel[eq] == level && invview(eq_to_diff)[eq] !== nothing,
-            varfilter = var->(var in level_vars && !(var in all_selected_states)))
-        for var in level_vars
-            if var_eq_matching[var] === unassigned
-                selected_state = invview(var_to_diff)[var]
-                push!(all_selected_states, selected_state)
-                #=
-                    # TODO: This is what the Matteson paper says, but it doesn't
-                    #       quite seem to work.
-                    while selected_state !== nothing
-                        push!(all_selected_states, selected_state)
-                        selected_state = invview(var_to_diff)[selected_state]
-                    end
-                =#
-            end
-        end
-        level += 1
-        level_vars = [var for var = 1:ndsts(graph) if varlevel[var] == level && invview(var_to_diff)[var] !== nothing]
-    end
+    var_eq_matching = pss_graph_modia!(graph, solvable_graph,
+        complete(var_eq_matching), var_to_diff, eq_to_diff, varlevel, inv_varlevel,
+        inv_eqlevel)
 
-    var_eq_matching = tear_graph_modia(graph, solvable_graph;
-        varfilter = var->!(var in all_selected_states))
-    var_eq_matching = Matching{Union{Unassigned, SelectedState}}(var_eq_matching)
-    for var in all_selected_states
-        var_eq_matching[var] = SelectedState()
-    end
-    return var_eq_matching, eq_to_diff
+    var_eq_matching, eq_to_diff
 end