Merge pull request #1347 from Keno/kf/graphscleanup1

BipartiteGraph: Clean up Graphs.jl integration

Author: YingboMa

src/bipartite_graph.jl

@@ -3,7 +3,7 @@ module BipartiteGraphs
 export BipartiteEdge, BipartiteGraph, DiCMOBiGraph, Unassigned, unassigned
 
 export 𝑠vertices, 𝑑vertices, has_𝑠vertex, has_𝑑vertex, 𝑠neighbors, 𝑑neighbors,
-       𝑠edges, 𝑑edges, nsrcs, ndsts, SRC, DST
+       𝑠edges, 𝑑edges, nsrcs, ndsts, SRC, DST, set_neighbors!
 
 using DocStringExtensions
 using UnPack
@@ -20,7 +20,7 @@ end
 ###
 ### Edges & Vertex
 ###
-@enum VertType SRC DST ALL
+@enum VertType SRC DST
 
 struct BipartiteEdge{I<:Integer} <: Graphs.AbstractEdge{I}
     src::I
@@ -189,10 +189,17 @@ function Graphs.add_vertex!(g::BipartiteGraph{T}, type::VertType) where T
     return true  # vertex successfully added
 end
 
+function set_neighbors!(g::BipartiteGraph, i::Integer, new_neighbors::AbstractVector)
+    old_nneighbors = length(g.fadjlist[i])
+    new_nneighbors = length(new_neighbors)
+    g.fadjlist[i] = new_neighbors
+    g.ne += new_nneighbors - old_nneighbors
+end
+
 ###
 ### Edges iteration
 ###
-Graphs.edges(g::BipartiteGraph) = BipartiteEdgeIter(g, Val(ALL))
+Graphs.edges(g::BipartiteGraph) = BipartiteEdgeIter(g, Val(SRC))
 𝑠edges(g::BipartiteGraph) = BipartiteEdgeIter(g, Val(SRC))
 𝑑edges(g::BipartiteGraph) = BipartiteEdgeIter(g, Val(DST))
 
@@ -202,8 +209,6 @@ struct BipartiteEdgeIter{T,G} <: Graphs.AbstractEdgeIter
 end
 
 Base.length(it::BipartiteEdgeIter) = ne(it.g)
-Base.length(it::BipartiteEdgeIter{ALL}) = 2ne(it.g)
-
 Base.eltype(it::BipartiteEdgeIter) = edgetype(it.g)
 
 function Base.iterate(it::BipartiteEdgeIter{SRC,<:BipartiteGraph{T}}, state=(1, 1, SRC)) where T
@@ -247,21 +252,6 @@ function Base.iterate(it::BipartiteEdgeIter{DST,<:BipartiteGraph{T}}, state=(1,
     return nothing
 end
 
-function Base.iterate(it::BipartiteEdgeIter{ALL,<:BipartiteGraph}, state=nothing)
-    if state === nothing
-        ss = iterate((@set it.type = Val(SRC)))
-    elseif state[3] === SRC
-        ss = iterate((@set it.type = Val(SRC)), state)
-    elseif state[3] == DST
-        ss = iterate((@set it.type = Val(DST)), state)
-    end
-    if ss === nothing && state[3] == SRC
-        return iterate((@set it.type = Val(DST)))
-    else
-        return ss
-    end
-end
-
 ###
 ### Utils
 ###
@@ -301,13 +291,20 @@ is acyclic if and only if the induced directed graph on the original bipartite
 graph is acyclic.
 
 """
-struct DiCMOBiGraph{Transposed, I, G<:BipartiteGraph{I}, M} <: Graphs.AbstractGraph{I}
+mutable struct DiCMOBiGraph{Transposed, I, G<:BipartiteGraph{I}, M} <: Graphs.AbstractGraph{I}
     graph::G
+    ne::Union{Missing, Int}
     matching::M
-    DiCMOBiGraph{Transposed}(g::G, m::M) where {Transposed, I, G<:BipartiteGraph{I}, M} =
-        new{Transposed, I, G, M}(g, 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)])
+end
+function DiCMOBiGraph{Transposed}(g::BipartiteGraph, m::M) where {Transposed, M}
+    DiCMOBiGraph{Transposed}(g, missing, m)
 end
-DiCMOBiGraph{Transposed}(g::BipartiteGraph) where {Transposed} = DiCMOBiGraph{Transposed}(g, Union{Unassigned, Int}[unassigned for i = 1:ndsts(g)])
+
 Graphs.is_directed(::Type{<:DiCMOBiGraph}) = true
 Graphs.nv(g::DiCMOBiGraph{Transposed}) where {Transposed} = Transposed ? ndsts(g.graph) : nsrcs(g.graph)
 Graphs.vertices(g::DiCMOBiGraph{Transposed}) where {Transposed} = Transposed ? 𝑑vertices(g.graph) : 𝑠vertices(g.graph)
@@ -318,6 +315,7 @@ struct CMONeighbors{Transposed, V}
     CMONeighbors{Transposed}(g::DiCMOBiGraph{Transposed}, v::V) where {Transposed, V} =
         new{Transposed, V}(g, v)
 end
+
 Graphs.outneighbors(g::DiCMOBiGraph{false}, v) = CMONeighbors{false}(g, v)
 Base.iterate(c::CMONeighbors{false}) = iterate(c, (c.g.graph.fadjlist[c.v],))
 function Base.iterate(c::CMONeighbors{false}, (l, state...))
@@ -336,12 +334,13 @@ function Base.iterate(c::CMONeighbors{false}, (l, state...))
     end
 end
 
+lift(f, x) = (x === unassigned || isnothing(x)) ? nothing : f(x)
+
+_vsrc(c::CMONeighbors{true}) = c.g.matching[c.v]
+_neighbors(c::CMONeighbors{true}) = lift(vsrc->c.g.graph.fadjlist[vsrc], _vsrc(c))
+Base.length(c::CMONeighbors{true}) = something(lift(length, _neighbors(c)), 1) - 1
 Graphs.inneighbors(g::DiCMOBiGraph{true}, v) = CMONeighbors{true}(g, v)
-function Base.iterate(c::CMONeighbors{true})
-    vsrc = c.g.matching[c.v]
-    vsrc === unassigned && return nothing
-    iterate(c, (c.g.graph.fadjlist[vsrc],))
-end
+Base.iterate(c::CMONeighbors{true}) = lift(ns->iterate(c, (ns,)), _neighbors(c))
 function Base.iterate(c::CMONeighbors{true}, (l, state...))
     while true
         r = iterate(l, state...)
@@ -354,4 +353,21 @@ function Base.iterate(c::CMONeighbors{true}, (l, state...))
     end
 end
 
+_edges(g::DiCMOBiGraph{Transposed}) where Transposed = Transposed ?
+    ((w=>v for w in inneighbors(g, v)) for v in vertices(g)) :
+    ((v=>w for w in outneighbors(g, v)) for v in vertices(g))
+_count(c::CMONeighbors{true}) = length(c)
+_count(c::CMONeighbors{false}) = count(_->true, c)
+
+Graphs.edges(g::DiCMOBiGraph) = (Graphs.SimpleEdge(p) for p in Iterators.flatten(_edges(g)))
+function Graphs.ne(g::DiCMOBiGraph)
+    if g.ne === missing
+        g.ne = mapreduce(x->_count(x.iter), +, _edges(g))
+    end
+    return g.ne
+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)
+
 end # module

src/structural_transformation/bipartite_tearing/modia_tearing.jl

@@ -35,6 +35,7 @@ function tearEquations!(ict::IncrementalCycleTracker, Gsolvable, es::Vector{Int}
             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 && break
             end

src/systems/alias_elimination.jl

@@ -127,6 +127,8 @@ struct AliasGraph <: AbstractDict{Int, Pair{Int, Int}}
     end
 end
 
+Base.length(ag::AliasGraph) = length(ag.eliminated)
+
 function Base.getindex(ag::AliasGraph, i::Integer)
     r = ag.aliasto[i]
     r === nothing && throw(KeyError(i))
@@ -281,7 +283,7 @@ function alias_eliminate_graph!(graph, varassoc, mm_orig::SparseMatrixCLIL)
 
     # Step 3: Reflect our update decitions back into the graph
     for (ei, e) in enumerate(mm.nzrows)
-        graph.fadjlist[e] = mm.row_cols[ei]
+        set_neighbors!(graph, e, mm.row_cols[ei])
     end
 
     return ag, mm

test/structural_transformation/utils.jl

@@ -27,9 +27,7 @@ sss = structure(sys)
 @test nv(solvable_graph) == 9 + 5
 @test varassoc == [0, 0, 0, 0, 1, 2, 3, 4, 0]
 
-se = collect(StructuralTransformations.𝑠edges(graph))
+se = collect(StructuralTransformations.edges(graph))
 @test se == mapreduce(vcat, enumerate(graph.fadjlist)) do (s, d)
     StructuralTransformations.BipartiteEdge.(s, d)
 end
-@test_throws ArgumentError collect(StructuralTransformations.𝑑edges(graph))
-@test_throws ArgumentError collect(StructuralTransformations.edges(graph))