Refactor (inv_)assign/{var/eq}assoc variables into proper structs

These various assignments and associations look very similar, but
I can never remember which one goes in which direction and which one
is a matching and which one is a derivative ladder. Make this all
more clear by factoring into to types `Matching` for matching and
`DiffGraph` for derivative ladders. Both have a vector structure
and the latter also exposes a graph structure on the same vertex
set as the BipartiteGraph, which is useful for plotting (not part
of this PR). I've also renamed the variables to make the direction
clear.

Author: Keno

src/bipartite_graph.jl

@@ -1,9 +1,11 @@
 module BipartiteGraphs
 
-export BipartiteEdge, BipartiteGraph, DiCMOBiGraph, Unassigned, unassigned
+export BipartiteEdge, BipartiteGraph, DiCMOBiGraph, Unassigned, unassigned,
+        Matching
 
 export 𝑠vertices, 𝑑vertices, has_𝑠vertex, has_𝑑vertex, 𝑠neighbors, 𝑑neighbors,
-       𝑠edges, 𝑑edges, nsrcs, ndsts, SRC, DST, set_neighbors!
+       𝑠edges, 𝑑edges, nsrcs, ndsts, SRC, DST, set_neighbors!, invview,
+       complete
 
 using DocStringExtensions
 using UnPack
@@ -17,6 +19,47 @@ struct Unassigned
     const unassigned = Unassigned.instance
 end
 
+struct Matching{V<:AbstractVector{<:Union{Unassigned, Int}}} <: AbstractVector{Union{Unassigned, Int}}
+    match::V
+    inv_match::Union{Nothing, V}
+end
+Matching(v::V) where {V<:AbstractVector{<:Union{Unassigned, Int}}} =
+    Matching{V}(v, nothing)
+Matching(m::Int) = Matching(Union{Int, Unassigned}[unassigned for _ = 1:m], nothing)
+Matching(m::Matching) = m
+
+Base.size(m::Matching) = Base.size(m.match)
+Base.getindex(m::Matching, i::Integer) = m.match[i]
+Base.iterate(m::Matching, state...) = iterate(m.match, state...)
+function Base.setindex!(m::Matching, v::Integer, i::Integer)
+    if m.inv_match !== nothing
+        m.inv_match[v] = i
+    end
+    return m.match[i] = v
+end
+
+function Base.push!(m::Matching, v::Union{Integer, Unassigned})
+    push!(m.match, v)
+    if v !== unassigned && m.inv_match !== nothing
+        m.inv_match[v] = length(m.match)
+    end
+end
+
+function complete(m::Matching)
+    m.inv_match !== nothing && return m
+    inv_match = Union{Unassigned, Int}[unassigned for _ = 1:length(m.match)]
+    for (i, eq) in enumerate(m.match)
+        eq === unassigned  && continue
+        inv_match[eq] = i
+    end
+    return Matching(collect(m.match), inv_match)
+end
+
+function invview(m::Matching)
+    m.inv_match === nothing && throw(ArgumentError("Backwards matching not defined. `complete` the matching first."))
+    return Matching(m.inv_match, m.match)
+end
+
 ###
 ### Edges & Vertex
 ###
@@ -291,15 +334,15 @@ is acyclic if and only if the induced directed graph on the original bipartite
 graph is acyclic.
 
 """
-mutable struct DiCMOBiGraph{Transposed, I, G<:BipartiteGraph{I}, M} <: Graphs.AbstractGraph{I}
+mutable struct DiCMOBiGraph{Transposed, I, G<:BipartiteGraph{I}, M <: Matching} <: Graphs.AbstractGraph{I}
     graph::G
     ne::Union{Missing, Int}
     matching::M
     DiCMOBiGraph{Transposed}(g::G, ne::Union{Missing, Int}, m::M) where {Transposed, I, G<:BipartiteGraph{I}, M} =
         new{Transposed, I, G, M}(g, ne, m)
 end
 function DiCMOBiGraph{Transposed}(g::BipartiteGraph) where {Transposed}
-    DiCMOBiGraph{Transposed}(g, 0, Union{Unassigned, Int}[unassigned for i = 1:ndsts(g)])
+    DiCMOBiGraph{Transposed}(g, 0, Matching(ndsts(g)))
 end
 function DiCMOBiGraph{Transposed}(g::BipartiteGraph, m::M) where {Transposed, M}
     DiCMOBiGraph{Transposed}(g, missing, m)

src/compat/incremental_cycles.jl

@@ -12,10 +12,10 @@ for a usage example.
 """
 abstract type IncrementalCycleTracker{I} <: AbstractGraph{I} end
 
-function (::Type{IncrementalCycleTracker})(s::AbstractGraph{I}; in_out_reverse=nothing) where {I}
+function (::Type{IncrementalCycleTracker})(s::AbstractGraph{I}; dir=:out) where {I}
     # TODO: Once we have more algorithms, the poly-algorithm decision goes here.
     # For now, we only have Algorithm N.
-    return DenseGraphICT_BFGT_N{something(in_out_reverse, false)}(s)
+    return DenseGraphICT_BFGT_N{something(dir == :in, false)}(s)
 end
 
 # Cycle Detection Interface
@@ -116,7 +116,7 @@ function Base.setindex!(vec::TransactionalVector, val, idx)
     return nothing
 end
 Base.getindex(vec::TransactionalVector, idx) = vec.v[idx]
-Base.size(vec) = size(vec.v)
+Base.size(vec::TransactionalVector) = size(vec.v)
 
 # Specific Algorithms
 

src/structural_transformation/StructuralTransformations.jl

@@ -22,6 +22,7 @@ using ModelingToolkit: ODESystem, AbstractSystem,var_from_nested_derivative, Dif
                        IncrementalCycleTracker, add_edge_checked!, topological_sort
 
 using ModelingToolkit.BipartiteGraphs
+import .BipartiteGraphs: invview
 using Graphs
 using ModelingToolkit.SystemStructures
 

src/structural_transformation/pantelides.jl

@@ -2,43 +2,39 @@
 ### Reassemble: structural information -> system
 ###
 
-function pantelides_reassemble(sys::ODESystem, eqassoc, assign)
+function pantelides_reassemble(sys::ODESystem, eq_to_diff, assign)
     s = structure(sys)
-    @unpack fullvars, varassoc = s
+    @unpack fullvars, var_to_diff = s
     # Step 1: write derivative equations
     in_eqs = equations(sys)
-    out_eqs = Vector{Any}(undef, length(eqassoc))
+    out_eqs = Vector{Any}(undef, nv(eq_to_diff))
     fill!(out_eqs, nothing)
     out_eqs[1:length(in_eqs)] .= in_eqs
 
-    out_vars = Vector{Any}(undef, length(varassoc))
+    out_vars = Vector{Any}(undef, nv(var_to_diff))
     fill!(out_vars, nothing)
     out_vars[1:length(fullvars)] .= fullvars
 
     D = Differential(get_iv(sys))
 
-    for (i, v) in enumerate(varassoc)
-        # fullvars[v] = D(fullvars[i])
-        v == 0 && continue
-        vi = out_vars[i]
+    for (varidx, diff) in edges(var_to_diff)
+        # fullvars[diff] = D(fullvars[var])
+        vi = out_vars[varidx]
         @assert vi !== nothing "Something went wrong on reconstructing states from variable association list"
         # `fullvars[i]` needs to be not a `D(...)`, because we want the DAE to be
         # first-order.
         if isdifferential(vi)
-            vi = out_vars[i] = diff2term(vi)
+            vi = out_vars[varidx] = diff2term(vi)
         end
-        out_vars[v] = D(vi)
+        out_vars[diff] = D(vi)
     end
 
     d_dict = Dict(zip(fullvars, 1:length(fullvars)))
     lhss = Set{Any}([x.lhs for x in in_eqs if isdiffeq(x)])
-    for (i, e) in enumerate(eqassoc)
-        if e === 0
-            continue
-        end
-        # LHS variable is looked up from varassoc
-        # the varassoc[i]-th variable is the differentiated version of var at i
-        eq = out_eqs[i]
+    for (eqidx, diff) in edges(eq_to_diff)
+        # LHS variable is looked up from var_to_diff
+        # the var_to_diff[i]-th variable is the differentiated version of var at i
+        eq = out_eqs[eqidx]
         lhs = if !(eq.lhs isa Symbolic)
             0
         elseif isdiffeq(eq)
@@ -58,7 +54,7 @@ function pantelides_reassemble(sys::ODESystem, eqassoc, assign)
         rhs = ModelingToolkit.expand_derivatives(D(eq.rhs))
         substitution_dict = Dict(x.lhs => x.rhs for x in out_eqs if x !== nothing && x.lhs isa Symbolic)
         sub_rhs = substitute(rhs, substitution_dict)
-        out_eqs[e] = lhs ~ sub_rhs
+        out_eqs[diff] = lhs ~ sub_rhs
     end
 
     final_vars = unique(filter(x->!(operation(x) isa Differential), fullvars))
@@ -78,16 +74,18 @@ Perform Pantelides algorithm.
 function pantelides!(sys::ODESystem; maxiters = 8000)
     s = structure(sys)
     # D(j) = assoc[j]
-    @unpack graph, fullvars, varassoc = s
-    iv = get_iv(sys)
+    @unpack graph, var_to_diff = s
+    return (sys, pantelides!(graph, var_to_diff)...)
+end
+
+function pantelides!(graph, var_to_diff; maxiters = 8000)
     neqs = nsrcs(graph)
-    nvars = length(varassoc)
+    nvars = nv(var_to_diff)
     vcolor = falses(nvars)
     ecolor = falses(neqs)
-    assign = Union{Unassigned, Int}[unassigned for _ = 1:nvars]
-    eqassoc = fill(0, neqs)
+    var_eq_matching = Matching(nvars)
+    eq_to_diff = DiffGraph(neqs)
     neqs′ = neqs
-    D = Differential(iv)
     for k in 1:neqs′
         eq′ = k
         pathfound = false
@@ -98,46 +96,46 @@ function pantelides!(sys::ODESystem; maxiters = 8000)
             #
             # the derivatives and algebraic variables are zeros in the variable
             # association list
-            varwhitelist = varassoc .== 0
+            varwhitelist = var_to_diff .== nothing
             resize!(vcolor, nvars)
             fill!(vcolor, false)
             resize!(ecolor, neqs)
             fill!(ecolor, false)
-            pathfound = find_augmenting_path(graph, eq′, assign, varwhitelist, vcolor, ecolor)
+            pathfound = find_augmenting_path(graph, eq′, var_eq_matching, varwhitelist, vcolor, ecolor)
             pathfound && break # terminating condition
             for var in eachindex(vcolor); vcolor[var] || continue
                 # introduce a new variable
                 nvars += 1
                 add_vertex!(graph, DST)
                 # the new variable is the derivative of `var`
-                varassoc[var] = nvars
-                push!(varassoc, 0)
-                push!(assign, unassigned)
+
+                add_edge!(var_to_diff, var, add_vertex!(var_to_diff))
+                push!(var_eq_matching, unassigned)
             end
 
             for eq in eachindex(ecolor); ecolor[eq] || continue
                 # introduce a new equation
                 neqs += 1
                 add_vertex!(graph, SRC)
                 # the new equation is created by differentiating `eq`
-                eqassoc[eq] = neqs
+                eq_diff = add_vertex!(eq_to_diff)
+                add_edge!(eq_to_diff, eq, eq_diff)
                 for var in 𝑠neighbors(graph, eq)
-                    add_edge!(graph, neqs, var)
-                    add_edge!(graph, neqs, varassoc[var])
+                    add_edge!(graph, eq_diff, var)
+                    add_edge!(graph, eq_diff, var_to_diff[var])
                 end
-                push!(eqassoc, 0)
             end
 
             for var in eachindex(vcolor); vcolor[var] || continue
                 # the newly introduced `var`s and `eq`s have the inherits
                 # assignment
-                assign[varassoc[var]] = eqassoc[assign[var]]
+                var_eq_matching[var_to_diff[var]] = eq_to_diff[var_eq_matching[var]]
             end
-            eq′ = eqassoc[eq′]
+            eq′ = eq_to_diff[eq′]
         end # for _ in 1:maxiters
         pathfound || error("maxiters=$maxiters reached! File a bug report if your system has a reasonable index (<100), and you are using the default `maxiters`. Try to increase the maxiters by `pantelides(sys::ODESystem; maxiters=1_000_000)` if your system has an incredibly high index and it is truly extremely large.")
     end # for k in 1:neqs′
-    return sys, assign, eqassoc
+    return var_eq_matching, eq_to_diff
 end
 
 """
@@ -150,6 +148,6 @@ instead, which calls this function internally.
 function dae_index_lowering(sys::ODESystem; kwargs...)
     s = get_structure(sys)
     (s isa SystemStructure) || (sys = initialize_system_structure(sys))
-    sys, assign, eqassoc = pantelides!(sys; kwargs...)
-    return pantelides_reassemble(sys, eqassoc, assign)
+    sys, var_eq_matching, eq_to_diff = pantelides!(sys; kwargs...)
+    return pantelides_reassemble(sys, eq_to_diff, var_eq_matching)
 end

src/structural_transformation/tearing.jl

@@ -7,13 +7,13 @@ instead, which calls this function internally.
 function tear_graph(sys)
     find_solvables!(sys)
     s = structure(sys)
-    @unpack graph, solvable_graph, assign, inv_assign, scc = s
+    @unpack graph, solvable_graph, var_eq_matching, scc = s
 
     @set! sys.structure.partitions = map(scc) do c
         ieqs = filter(eq->isalgeq(s, eq), c)
-        vars = inv_assign[ieqs]
+        vars = Int[var for var in invview(var_eq_matching)[ieqs] if var !== unassigned]
 
-        ict = IncrementalCycleTracker(DiCMOBiGraph{true}(graph); in_out_reverse=true)
+        ict = IncrementalCycleTracker(DiCMOBiGraph{true}(graph); dir=:in)
         SystemPartition(tearEquations!(ict, solvable_graph.fadjlist, ieqs, vars)...)
     end
     return sys
@@ -26,7 +26,7 @@ end
 
 function tearing_reassemble(sys; simplify=false)
     s = structure(sys)
-    @unpack fullvars, partitions, assign, inv_assign, graph, scc = s
+    @unpack fullvars, partitions, var_eq_matching, graph, scc = s
     eqs = equations(sys)
 
     ### extract partition information
@@ -152,7 +152,7 @@ function tearing_reassemble(sys; simplify=false)
                 if abs(rhs) > 100eps(float(rhs))
                     @warn "The equation $eq is not consistent. It simplifed to 0 == $rhs."
                 end
-                neweqs[ridx] = 0 ~ fullvars[inv_assign[ieq]]
+                neweqs[ridx] = 0 ~ fullvars[invview(var_eq_matching)[ieq]]
             end
         end
     end
@@ -205,14 +205,10 @@ function algebraic_equations_scc(sys)
 
     # skip over differential equations
     algvars = isalgvar.(Ref(s), 1:ndsts(s.graph))
-    eqs = equations(sys)
-    assign = matching(s, algvars, s.algeqs)
-
-    components = find_scc(s.graph, assign)
-    inv_assign = inverse_mapping(assign)
+    var_eq_matching = complete(matching(s, algvars, s.algeqs))
+    components = find_scc(s.graph, var_eq_matching)
 
-    @set! sys.structure.assign = assign
-    @set! sys.structure.inv_assign = inv_assign
+    @set! sys.structure.var_eq_matching = var_eq_matching
     @set! sys.structure.scc = components
     return sys
 end