Add an attribute has_loops in AdjacencyGraph (#186)

* Add an attribute has_loops in AdjacencyGraph

* Rename `has_loops` into `has_diagonal`, make it statically inferred

* Fixes

---------

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

Author: amontoison

src/coloring.jl

@@ -557,9 +557,11 @@ function postprocess!(
     color_used = zeros(Bool, maximum(color))
 
     # nonzero diagonal coefficients force the use of their respective color (there can be no neutral colors if the diagonal is fully nonzero)
-    for i in axes(S, 1)
-        if !iszero(S[i, i])
-            color_used[color[i]] = true
+    if has_diagonal(g)
+        for i in axes(S, 1)
+            if !iszero(S[i, i])
+                color_used[color[i]] = true
+            end
         end
     end
 

src/decompression.jl

@@ -415,16 +415,22 @@ end
 function decompress!(
     A::AbstractMatrix, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F
 )
-    (; color, star_set) = result
+    (; ag, color, star_set) = result
     (; star, hub, spokes) = star_set
-    S = result.ag.S
+    (; S) = ag
     uplo == :F && check_same_pattern(A, S)
     fill!(A, zero(eltype(A)))
-    for i in axes(A, 1)
-        if !iszero(S[i, i])
-            A[i, i] = B[i, color[i]]
+
+    # Recover the diagonal coefficients of A
+    if has_diagonal(ag)
+        for i in axes(A, 1)
+            if !iszero(S[i, i])
+                A[i, i] = B[i, color[i]]
+            end
         end
     end
+
+    # Recover the off-diagonal coefficients of A
     for s in eachindex(hub, spokes)
         j = abs(hub[s])
         cj = color[j]
@@ -447,15 +453,21 @@ function decompress_single_color!(
     result::StarSetColoringResult,
     uplo::Symbol=:F,
 )
-    (; color, group, star_set) = result
+    (; ag, color, group, star_set) = result
     (; hub, spokes) = star_set
-    S = result.ag.S
+    (; S) = ag
     uplo == :F && check_same_pattern(A, S)
-    for i in axes(A, 1)
-        if !iszero(S[i, i]) && color[i] == c
-            A[i, i] = b[i]
+
+    # Recover the diagonal coefficients of A
+    if has_diagonal(ag)
+        for i in axes(A, 1)
+            if !iszero(S[i, i]) && color[i] == c
+                A[i, i] = b[i]
+            end
         end
     end
+
+    # Recover the off-diagonal coefficients of A
     for s in eachindex(hub, spokes)
         j = abs(hub[s])
         if color[j] == c
@@ -475,8 +487,8 @@ end
 function decompress!(
     A::SparseMatrixCSC, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F
 )
-    (; compressed_indices) = result
-    S = result.ag.S
+    (; ag, compressed_indices) = result
+    (; S) = ag
     nzA = nonzeros(A)
     if uplo == :F
         check_same_pattern(A, S)
@@ -505,8 +517,8 @@ end
 function decompress!(
     A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
 )
-    (; color, vertices_by_tree, reverse_bfs_orders, buffer) = result
-    S = result.ag.S
+    (; ag, color, vertices_by_tree, reverse_bfs_orders, buffer) = result
+    (; S) = ag
     uplo == :F && check_same_pattern(A, S)
     R = eltype(A)
     fill!(A, zero(R))
@@ -518,9 +530,11 @@ function decompress!(
     end
 
     # Recover the diagonal coefficients of A
-    for i in axes(A, 1)
-        if !iszero(S[i, i])
-            A[i, i] = B[i, color[i]]
+    if has_diagonal(ag)
+        for i in axes(A, 1)
+            if !iszero(S[i, i])
+                A[i, i] = B[i, color[i]]
+            end
         end
     end
 
@@ -551,6 +565,7 @@ function decompress!(
     uplo::Symbol=:F,
 ) where {R<:Real}
     (;
+        ag,
         color,
         vertices_by_tree,
         reverse_bfs_orders,
@@ -560,7 +575,7 @@ function decompress!(
         upper_triangle_offsets,
         buffer,
     ) = result
-    S = result.ag.S
+    (; S) = ag
     A_colptr = A.colptr
     nzA = nonzeros(A)
     uplo == :F && check_same_pattern(A, S)
@@ -572,22 +587,24 @@ function decompress!(
     end
 
     # Recover the diagonal coefficients of A
-    if uplo == :L
-        for i in diagonal_indices
-            # A[i, i] is the first element in column i
-            nzind = A_colptr[i]
-            nzA[nzind] = B[i, color[i]]
-        end
-    elseif uplo == :U
-        for i in diagonal_indices
-            # A[i, i] is the last element in column i
-            nzind = A_colptr[i + 1] - 1
-            nzA[nzind] = B[i, color[i]]
-        end
-    else  # uplo == :F
-        for (k, i) in enumerate(diagonal_indices)
-            nzind = diagonal_nzind[k]
-            nzA[nzind] = B[i, color[i]]
+    if has_diagonal(ag)
+        if uplo == :L
+            for i in diagonal_indices
+                # A[i, i] is the first element in column i
+                nzind = A_colptr[i]
+                nzA[nzind] = B[i, color[i]]
+            end
+        elseif uplo == :U
+            for i in diagonal_indices
+                # A[i, i] is the last element in column i
+                nzind = A_colptr[i + 1] - 1
+                nzA[nzind] = B[i, color[i]]
+            end
+        else  # uplo == :F
+            for (k, i) in enumerate(diagonal_indices)
+                nzind = diagonal_nzind[k]
+                nzA[nzind] = B[i, color[i]]
+            end
         end
     end
 

src/graph.jl

@@ -96,7 +96,7 @@ end
 ## Adjacency graph
 
 """
-    AdjacencyGraph{T}
+    AdjacencyGraph{T,has_diagonal}
 
 Undirected graph without self-loops representing the nonzeros of a symmetric matrix (typically a Hessian matrix).
 
@@ -111,27 +111,44 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 
 # Fields
 
-- `S::SparsityPatternCSC{T}`
+- `S::SparsityPatternCSC{T}`: Underlying sparsity pattern, whose diagonal is empty whenever `has_diagonal` is `false`
 
 # 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}
+struct AdjacencyGraph{T,has_diagonal}
     S::SparsityPatternCSC{T}
 end
 
-AdjacencyGraph(A::AbstractMatrix) = AdjacencyGraph(SparseMatrixCSC(A))
-AdjacencyGraph(A::SparseMatrixCSC) = AdjacencyGraph(SparsityPatternCSC(A))
+function AdjacencyGraph(S::SparsityPatternCSC{T}; has_diagonal::Bool=true) where {T}
+    return AdjacencyGraph{T,has_diagonal}(S)
+end
+
+function AdjacencyGraph(A::SparseMatrixCSC; has_diagonal::Bool=true)
+    return AdjacencyGraph(SparsityPatternCSC(A); has_diagonal)
+end
+
+function AdjacencyGraph(A::AbstractMatrix; has_diagonal::Bool=true)
+    return AdjacencyGraph(SparseMatrixCSC(A); has_diagonal)
+end
 
 pattern(g::AdjacencyGraph) = g.S
 nb_vertices(g::AdjacencyGraph) = pattern(g).n
 vertices(g::AdjacencyGraph) = 1:nb_vertices(g)
 
-function neighbors(g::AdjacencyGraph, v::Integer)
+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}
     S = pattern(g)
-    neighbors_with_loops = view(rowvals(S), nzrange(S, v))
-    return Iterators.filter(!=(v), neighbors_with_loops)  # TODO: optimize
+    neighbors_v = view(rowvals(S), nzrange(S, v))
+    return neighbors_v
 end
 
 function degree(g::AdjacencyGraph, v::Integer)

src/interface.jl

@@ -245,7 +245,7 @@ function coloring(
         spzeros(T, n, n) SparseMatrixCSC(Aᵀ)
         SparseMatrixCSC(A) spzeros(T, m, m)
     ]  # TODO: slow
-    ag = AdjacencyGraph(A_and_Aᵀ)
+    ag = AdjacencyGraph(A_and_Aᵀ; has_diagonal=false)
     if decompression == :direct
         color, star_set = star_coloring(ag, algo.order; postprocessing=algo.postprocessing)
         symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)

src/result.jl

@@ -295,23 +295,25 @@ function TreeSetColoringResult(
     decompression_eltype::Type{R},
 ) where {R}
     (; vertices_by_tree, reverse_bfs_orders) = tree_set
-    S = ag.S
+    (; S) = ag
     nvertices = length(color)
     group = group_by_color(color)
+    rv = rowvals(S)
 
     # Vector for the decompression of the diagonal coefficients
     diagonal_indices = Int[]
     diagonal_nzind = Int[]
     ndiag = 0
 
-    rv = rowvals(S)
-    for j in axes(S, 2)
-        for k in nzrange(S, j)
-            i = rv[k]
-            if i == j
-                push!(diagonal_indices, i)
-                push!(diagonal_nzind, k)
-                ndiag += 1
+    if has_diagonal(ag)
+        for j in axes(S, 2)
+            for k in nzrange(S, j)
+                i = rv[k]
+                if i == j
+                    push!(diagonal_indices, i)
+                    push!(diagonal_nzind, k)
+                    ndiag += 1
+                end
             end
         end
     end
@@ -497,7 +499,7 @@ struct BicoloringResult{
 } <: AbstractColoringResult{:nonsymmetric,:bidirectional,decompression}
     "matrix that was colored"
     A::M
-    "adjacency graph that was used for coloring (constructed from the bipartite graph)"
+    "augmented adjacency graph that was used for bicoloring"
     abg::G
     "one integer color for each column"
     column_color::Vector{Int}
@@ -507,7 +509,7 @@ struct BicoloringResult{
     column_group::V
     "color groups for rows"
     row_group::V
-    "result for the coloring of the symmetric 2x2 block matrix"
+    "result for the coloring of the symmetric 2 x 2 block matrix"
     symmetric_result::SR
     "column color to index"
     col_color_ind::Dict{Int,Int}