Merge pull request #1350 from Keno/kf/assignstructs

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

Author: YingboMa

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