Backport the modifications of Guillaume

amontoison · amontoison · commit 750587672df8 · 2025-03-31T19:06:22.000-05:00
diff --git a/src/coloring.jl b/src/coloring.jl
@@ -82,18 +82,18 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessing::
     nv = nb_vertices(g)
     ne = nb_edges(g)
     color = zeros(Int, nv)
-    forbidden_colors = zeros(Int, nv)
-    edge_to_index = build_edge_to_index(S, forbidden_colors)
+    forbidden_colors = g.vertex_buffer  # Vector of length nv
+    fill!(forbidden_colors, 0)
     first_neighbor = fill((0, 0, 0), nv)  # at first no neighbors have been encountered
     treated = zeros(Int, nv)
     star = Vector{Int}(undef, ne)
     hub = Int[]  # one hub for each star, including the trivial ones
     vertices_in_order = vertices(g, order)
 
     for v in vertices_in_order
-        for (iw, w) in enumerate(neighbors2(g, v))
-            (v == w || iszero(color[w])) && continue
-            index_vw = edge_to_index[S.colptr[v] + iw - 1]
+        for (w, index_vw) in neighbors_with_edge_indices(g, v)
+            !has_diagonal(g) || (v == w && continue)
+            iszero(color[w]) && continue
             forbidden_colors[color[w]] = v
             (p, q, _) = first_neighbor[color[w]]
             if p == v  # Case 1
@@ -105,9 +105,9 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessing::
                 _treat!(treated, forbidden_colors, g, v, w, color)
             else
                 first_neighbor[color[w]] = (v, w, index_vw)
-                for (ix, x) in enumerate(neighbors2(g, w))
-                    (w == x || x == v || iszero(color[x])) && continue
-                    index_wx = edge_to_index[S.colptr[w] + ix - 1]
+                for (x, index_wx) in neighbors_with_edge_indices(g, w)
+                    !has_diagonal(g) || (w == x && continue)
+                    (x == v || iszero(color[x])) && continue
                     if x == hub[star[index_wx]]  # potential Case 2 (which is always false for trivial stars with two vertices, since the associated hub is negative)
                         forbidden_colors[color[x]] = v
                     end
@@ -120,14 +120,13 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessing::
                 break
             end
         end
-        _update_stars!(star, hub, g, v, color, first_neighbor, edge_to_index)
+        _update_stars!(star, hub, g, v, color, first_neighbor)
     end
     star_set = StarSet(star, hub)
     if postprocessing
-        # Reuse the vector forbidden_colors to compute offsets during post-processing
-        postprocess!(color, star_set, g, forbidden_colors, edge_to_index)
+        postprocess!(color, star_set, g)
     end
-    return color, star_set, edge_to_index
+    return color, star_set
 end
 
 function _treat!(
@@ -141,6 +140,7 @@ function _treat!(
     color::AbstractVector{<:Integer},
 )
     for x in neighbors(g, w)
+        !has_diagonal(g) || (w == x && continue)
         iszero(color[x]) && continue
         forbidden_colors[color[x]] = v
     end
@@ -157,17 +157,15 @@ function _update_stars!(
     v::Integer,
     color::AbstractVector{<:Integer},
     first_neighbor::AbstractVector{<:Tuple},
-    edge_to_index::AbstractVector{<:Integer},
 )
     S = pattern(g)
-    for (iw, w) in enumerate(neighbors2(g, v))
-        (v == w || iszero(color[w])) && continue
-        index_vw = edge_to_index[S.colptr[v] + iw - 1]
+    for (w, index_vw) in neighbors_with_edge_indices(g, v)
+        !has_diagonal(g) || (v == w && continue)
+        iszero(color[w]) && continue
         x_exists = false
-        for (ix, x) in enumerate(neighbors2(g, w))
-            (w == x) && continue
+        for (x, index_wx) in neighbors_with_edge_indices(g, w)
+            !has_diagonal(g) || (w == x && continue)
             if x != v && color[x] == color[v]  # vw, wx ∈ E
-                index_wx = edge_to_index[S.colptr[w] + ix - 1]
                 star_wx = star[index_wx]
                 hub[star_wx] = w  # this may already be true
                 star[index_vw] = star_wx
@@ -235,24 +233,26 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessin
     nv = nb_vertices(g)
     ne = nb_edges(g)
     color = zeros(Int, nv)
-    forbidden_colors = zeros(Int, nv)
-    edge_to_index = build_edge_to_index(S, forbidden_colors)
+    forbidden_colors = g.vertex_buffer  # Vector of length nv
+    fill!(forbidden_colors, 0)
     first_neighbor = fill((0, 0, 0), nv)  # at first no neighbors have been encountered
     first_visit_to_tree = fill((0, 0), ne)
     forest = Forest{Int}(ne)
     vertices_in_order = vertices(g, order)
 
     for v in vertices_in_order
         for w in neighbors(g, v)
+            !has_diagonal(g) || (v == w && continue)
             iszero(color[w]) && continue
             forbidden_colors[color[w]] = v
         end
         for w in neighbors(g, v)
+            !has_diagonal(g) || (v == w && continue)
             iszero(color[w]) && continue
-            for (ix, x) in enumerate(neighbors2(g, w))
-                (w == x || iszero(color[x])) && continue
+            for (x, index_wx) in neighbors_with_edge_indices(g, w)
+                !has_diagonal(g) || (w == x && continue)
+                iszero(color[x]) && continue
                 if forbidden_colors[color[x]] != v
-                    index_wx = edge_to_index[S.colptr[w] + ix - 1]
                     _prevent_cycle!(
                         v,
                         w,
@@ -272,28 +272,27 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessin
                 break
             end
         end
-        for (iw, w) in enumerate(neighbors2(g, v))  # grow two-colored stars around the vertex v
-            (v == w || iszero(color[w])) && continue
-            index_vw = edge_to_index[S.colptr[v] + iw - 1]
+        for (w, index_vw) in neighbors_with_edge_indices(g, v)  # grow two-colored stars around the vertex v
+            !has_diagonal(g) || (v == w && continue)
+            iszero(color[w]) && continue
             _grow_star!(v, w, index_vw, color, first_neighbor, forest)
         end
-        for (iw, w) in enumerate(neighbors2(g, v))
-            (v == w || iszero(color[w])) && continue
-            for (ix, x) in enumerate(neighbors2(g, w))
-                (w == x || x == v || iszero(color[x])) && continue
+        for (w, index_vw) in neighbors_with_edge_indices(g, v)
+            !has_diagonal(g) || (v == w && continue)
+            iszero(color[w]) && continue
+            for (x, index_wx) in neighbors_with_edge_indices(g, w)
+                !has_diagonal(g) || (w == x && continue)
+                (x == v || iszero(color[x])) && continue
                 if color[x] == color[v]
-                    index_vw = edge_to_index[S.colptr[v] + iw - 1]
-                    index_wx = edge_to_index[S.colptr[w] + ix - 1]
                     _merge_trees!(v, w, x, index_vw, index_wx, forest)  # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂
                 end
             end
         end
     end
 
-    tree_set = TreeSet(g, forest, edge_to_index)
+    tree_set = TreeSet(g, forest)
     if postprocessing
-        # Reuse the vector forbidden_colors to compute offsets during post-processing
-        postprocess!(color, tree_set, g, forbidden_colors, edge_to_index)
+        postprocess!(color, tree_set, g)
     end
     return color, tree_set
 end
@@ -374,8 +373,9 @@ struct TreeSet
     is_star::Vector{Bool}
 end
 
-function TreeSet(g::AdjacencyGraph, forest::Forest{Int}, edge_to_index::Vector{Int})
+function TreeSet(g::AdjacencyGraph, forest::Forest{Int})
     S = pattern(g)
+    edge_to_index = edge_indices(g)
     nv = nb_vertices(g)
     nt = forest.num_trees
 
@@ -520,10 +520,9 @@ function postprocess!(
     color::AbstractVector{<:Integer},
     star_or_tree_set::Union{StarSet,TreeSet},
     g::AdjacencyGraph,
-    offsets::Vector{Int},
-    edge_to_index::Vector{Int},
 )
-    (; S) = g
+    S = pattern(g)
+    edge_to_index = edge_indices(g)
     # flag which colors are actually used during decompression
     nb_colors = maximum(color)
     color_used = zeros(Bool, nb_colors)
@@ -626,6 +625,7 @@ function postprocess!(
 
     # if at least one of the colors is useless, modify the color assignments of vertices
     if any(!, color_used)
+        offsets = g.vertex_buffer
         num_colors_useless = 0
 
         # determine what are the useless colors and compute the offsets
diff --git a/src/graph.jl b/src/graph.jl
@@ -36,7 +36,7 @@ SparseArrays.nzrange(S::SparsityPatternCSC, j::Integer) = S.colptr[j]:(S.colptr[
 
 Return a [`SparsityPatternCSC`](@ref) corresponding to the transpose of `S`.
 """
-function Base.transpose(S::SparsityPatternCSC{T}) where {T}
+function Base.transpose(S::SparsityPatternCSC{T}) where {T<:Integer}
     m, n = size(S)
     nnzA = nnz(S)
     A_colptr = S.colptr
@@ -98,10 +98,10 @@ end
 
 Return a [`SparsityPatternCSC`](@ref) corresponding to the matrix `[0 Aᵀ; A 0]`, with a minimum of allocations.
 """
-bidirectional_pattern(A::AbstractMatrix; symmetric_pattern) =
+bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool) =
     bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
 
-function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern)
+function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern::Bool)
     m, n = size(S)
     p = m + n
     nnzS = nnz(S)
@@ -165,11 +165,10 @@ end
 
 # Note that we don't need to do that for bicoloring,
 # we can build that in the same time than the transposed sparsity pattern of A
-function build_edge_to_index(S::SparsityPatternCSC, forbidden_colors::Vector{Int})
+function build_edge_to_index(S::SparsityPatternCSC{T}) where {T<:Integer}
     # edge_to_index gives an index for each edge
-    edge_to_index = Vector{Int}(undef, nnz(S))
-    # use forbidden_colors (or color) for the offsets of each column
-    offsets = forbidden_colors
+    edge_to_index = Vector{T}(undef, nnz(S))
+    offsets = zeros(T, S.n)
     counter = 0
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -187,8 +186,7 @@ function build_edge_to_index(S::SparsityPatternCSC, forbidden_colors::Vector{Int
             end
         end
     end
-    fill!(offsets, 0)
-    return edge_to_index
+    return edge_to_index, offsets
 end
 
 ## Adjacency graph
@@ -205,22 +203,27 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 
 # Constructors
 
-    AdjacencyGraph(A::SparseMatrixCSC)
+    AdjacencyGraph(A::SparseMatrixCSC; has_diagonal::Bool=true)
 
 # Fields
 
 - `S::SparsityPatternCSC{T}`: Underlying sparsity pattern, whose diagonal is empty whenever `has_diagonal` is `false`
+- `edge_to_index::Vector{T}`: A vector mapping each nonzero of `S` to a unique edge index (ignoring diagonal and accounting for symmetry, so that `(i, j)` and `(j, i)` get the same index)
+- `vertex_buffer::Vector{T}`: A buffer of length `S.n`, which is the number of vertices in the graph
 
 # References
 
 > [_What Color Is Your Jacobian? SparsityPatternCSC Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 struct AdjacencyGraph{T,has_diagonal}
     S::SparsityPatternCSC{T}
+    edge_to_index::Vector{T}
+    vertex_buffer::Vector{T}
 end
 
 function AdjacencyGraph(S::SparsityPatternCSC{T}; has_diagonal::Bool=true) where {T}
-    return AdjacencyGraph{T,has_diagonal}(S)
+    edge_to_index, vertex_buffer = build_edge_to_index(S)
+    return AdjacencyGraph{T,has_diagonal}(S, edge_to_index, vertex_buffer)
 end
 
 function AdjacencyGraph(A::SparseMatrixCSC; has_diagonal::Bool=true)
@@ -232,30 +235,27 @@ function AdjacencyGraph(A::AbstractMatrix; has_diagonal::Bool=true)
 end
 
 pattern(g::AdjacencyGraph) = g.S
+edge_indices(g::AdjacencyGraph) = g.edge_to_index
 nb_vertices(g::AdjacencyGraph) = pattern(g).n
 vertices(g::AdjacencyGraph) = 1:nb_vertices(g)
+has_diagonal(::AdjacencyGraph{T,hd}) where {T,hd} = hd
 
-has_diagonal(::AdjacencyGraph{T,hl}) where {T,hl} = hl
-
-function neighbors(g::AdjacencyGraph{T,true}, v::Integer) where {T}
-    S = pattern(g)
-    neighbors_v = view(rowvals(S), nzrange(S, v))
-    return Iterators.filter(!=(v), neighbors_v)  # TODO: optimize
-end
-
-function neighbors(g::AdjacencyGraph{T,false}, v::Integer) where {T}
+function neighbors(g::AdjacencyGraph, v::Integer)
     S = pattern(g)
     neighbors_v = view(rowvals(S), nzrange(S, v))
     return neighbors_v
 end
 
-function neighbors2(g::AdjacencyGraph, v::Integer)
+function neighbors_with_edge_indices(g::AdjacencyGraph, v::Integer)
     S = pattern(g)
     neighbors_v = view(rowvals(S), nzrange(S, v))
-    return neighbors_v
+    edges_indices_v = view(edge_indices(g), nzrange(S, v))
+    return zip(neighbors_v, edges_indices_v)
 end
 
-function degree(g::AdjacencyGraph, v::Integer)
+degree(g::AdjacencyGraph{T,false}, v::Integer) where {T} = g.S.colptr[v + 1] - g.S.colptr[v]
+
+function degree(g::AdjacencyGraph{T,true}, v::Integer) where {T}
     d = 0
     for u in neighbors(g, v)
         if u != v
@@ -265,12 +265,12 @@ function degree(g::AdjacencyGraph, v::Integer)
     return d
 end
 
-function nb_edges(g::AdjacencyGraph)
+nb_edges(g::AdjacencyGraph{T,false}) where {T} = nnz(g.S) ÷ 2
+
+function nb_edges(g::AdjacencyGraph{T,true}) where {T}
     ne = 0
     for v in vertices(g)
-        for u in neighbors(g, v)
-            ne += 1
-        end
+        ne += degree(g, v)
     end
     return ne ÷ 2
 end
@@ -280,6 +280,7 @@ minimum_degree(g::AdjacencyGraph) = minimum(Base.Fix1(degree, g), vertices(g))
 
 function has_neighbor(g::AdjacencyGraph, v::Integer, u::Integer)
     for w in neighbors(g, v)
+        !has_diagonal(g) || (v == w && continue)
         if w == u
             return true
         end
@@ -314,7 +315,7 @@ A `BipartiteGraph` has two sets of vertices, one for the rows of `A` (which we c
 
 # Constructors
 
-    BipartiteGraph(A::SparseMatrixCSC; symmetric_pattern=false)
+    BipartiteGraph(A::SparseMatrixCSC; symmetric_pattern::Bool=false)
 
 When `symmetric_pattern` is `true`, this construction is more efficient.
 
diff --git a/src/interface.jl b/src/interface.jl
@@ -266,9 +266,9 @@ function _coloring(
     symmetric_pattern::Bool,
 )
     ag = AdjacencyGraph(A)
-    color, star_set, edge_to_index = star_coloring(ag, algo.order, algo.postprocessing)
+    color, star_set = star_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
-        return StarSetColoringResult(A, ag, color, edge_to_index, star_set)
+        return StarSetColoringResult(A, ag, color, star_set)
     else
         return color
     end
@@ -301,11 +301,9 @@ function _coloring(
 ) where {R}
     A_and_Aᵀ = bidirectional_pattern(A; symmetric_pattern)
     ag = AdjacencyGraph(A_and_Aᵀ; has_diagonal=false)
-    color, star_set, edge_to_index = star_coloring(ag, algo.order, algo.postprocessing)
+    color, star_set = star_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
-        symmetric_result = StarSetColoringResult(
-            A_and_Aᵀ, ag, color, edge_to_index, star_set
-        )
+        symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
         return BicoloringResult(A, ag, symmetric_result, R)
     else
         row_color, column_color, _ = remap_colors(color, maximum(color), size(A)...)
diff --git a/src/order.jl b/src/order.jl
@@ -262,6 +262,7 @@ function vertices(
         u = pop_next_candidate!(db; direction)
         direction == :low2high ? push!(π, u) : pushfirst!(π, u)
         for v in neighbors(g, u)
+            !has_diagonal(g) || (u == v && continue)
             already_ordered(db, v) && continue
             update_bucket!(db, v; degtype, direction)
         end
diff --git a/src/result.jl b/src/result.jl
@@ -239,14 +239,10 @@ struct StarSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V} <:
 end
 
 function StarSetColoringResult(
-    A::AbstractMatrix,
-    ag::AdjacencyGraph,
-    color::Vector{Int},
-    edge_to_index::Vector{Int},
-    star_set::StarSet,
+    A::AbstractMatrix, ag::AdjacencyGraph, color::Vector{Int}, star_set::StarSet
 )
-    # Reuse edge_to_index to store the compressed indices for decompression
-    compressed_indices = edge_to_index
+    # Reuse ag.edge_to_index to store the compressed indices for decompression
+    compressed_indices = edge_indices(ag)
 
     (; star, hub) = star_set
     S = pattern(ag)
