Add nb_self_loops in AdjacencyGraph

Author: amontoison

docs/src/dev.md

@@ -58,7 +58,7 @@ SparseMatrixColorings.valid_dynamic_order
 ```@docs
 SparseMatrixColorings.respectful_similar
 SparseMatrixColorings.matrix_versions
-SparseMatrixColorings.same_pattern
+SparseMatrixColorings.compatible_pattern
 ```
 
 ## Visualization

src/decompression.jl

@@ -339,9 +339,9 @@ end
 ## ColumnColoringResult
 
 function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::ColumnColoringResult)
-    (; color) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, color) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     fill!(A, zero(eltype(A)))
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -357,9 +357,9 @@ end
 function decompress_single_color!(
     A::AbstractMatrix, b::AbstractVector, c::Integer, result::ColumnColoringResult
 )
-    (; group) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, group) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     rvS = rowvals(S)
     for j in group[c]
         for k in nzrange(S, j)
@@ -371,9 +371,9 @@ function decompress_single_color!(
 end
 
 function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::ColumnColoringResult)
-    (; compressed_indices) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, compressed_indices) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     nzA = nonzeros(A)
     for k in eachindex(nzA, compressed_indices)
         nzA[k] = B[compressed_indices[k]]
@@ -384,9 +384,9 @@ end
 function decompress_single_color!(
     A::SparseMatrixCSC, b::AbstractVector, c::Integer, result::ColumnColoringResult
 )
-    (; group) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, group) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     rvS = rowvals(S)
     nzA = nonzeros(A)
     for j in group[c]
@@ -401,9 +401,9 @@ end
 ## RowColoringResult
 
 function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::RowColoringResult)
-    (; color) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, color) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     fill!(A, zero(eltype(A)))
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -419,9 +419,9 @@ end
 function decompress_single_color!(
     A::AbstractMatrix, b::AbstractVector, c::Integer, result::RowColoringResult
 )
-    (; group) = result
-    S, Sᵀ = result.bg.S2, result.bg.S1
-    check_same_pattern(A, S)
+    (; bg, group) = result
+    S, Sᵀ = bg.S2, bg.S1
+    check_compatible_pattern(A, bg)
     rvSᵀ = rowvals(Sᵀ)
     for i in group[c]
         for k in nzrange(Sᵀ, i)
@@ -433,9 +433,9 @@ function decompress_single_color!(
 end
 
 function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::RowColoringResult)
-    (; compressed_indices) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, compressed_indices) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     nzA = nonzeros(A)
     for k in eachindex(nzA, compressed_indices)
         nzA[k] = B[compressed_indices[k]]
@@ -450,7 +450,7 @@ function decompress!(
 )
     (; ag, compressed_indices) = result
     (; S) = ag
-    uplo == :F && check_same_pattern(A, S)
+    check_compatible_pattern(A, ag, uplo)
     fill!(A, zero(eltype(A)))
 
     rvS = rowvals(S)
@@ -474,7 +474,7 @@ function decompress_single_color!(
 )
     (; ag, compressed_indices, group) = result
     (; S) = ag
-    uplo == :F && check_same_pattern(A, S)
+    check_compatible_pattern(A, ag, uplo)
 
     lower_index = (c - 1) * S.n + 1
     upper_index = c * S.n
@@ -508,8 +508,8 @@ function decompress!(
     (; ag, compressed_indices) = result
     (; S) = ag
     nzA = nonzeros(A)
+    check_compatible_pattern(A, ag, uplo)
     if uplo == :F
-        check_same_pattern(A, S)
         for k in eachindex(nzA, compressed_indices)
             nzA[k] = B[compressed_indices[k]]
         end
@@ -537,7 +537,7 @@ function decompress!(
     (; ag, color, reverse_bfs_orders, tree_edge_indices, nt, diagonal_indices, buffer) =
         result
     (; S) = ag
-    uplo == :F && check_same_pattern(A, S)
+    check_compatible_pattern(A, ag, uplo)
     R = eltype(A)
     fill!(A, zero(R))
 
@@ -606,7 +606,7 @@ function decompress!(
     (; S) = ag
     A_colptr = A.colptr
     nzA = nonzeros(A)
-    uplo == :F && check_same_pattern(A, S)
+    check_compatible_pattern(A, ag, uplo)
 
     if eltype(buffer) == R
         buffer_right_type = buffer
@@ -707,9 +707,10 @@ function decompress!(
     result::LinearSystemColoringResult,
     uplo::Symbol=:F,
 )
-    (; color, strict_upper_nonzero_inds, M_factorization, strict_upper_nonzeros_A) = result
-    S = result.ag.S
-    uplo == :F && check_same_pattern(A, S)
+    (; ag, color, strict_upper_nonzero_inds, M_factorization, strict_upper_nonzeros_A) =
+        result
+    S = ag.S
+    check_compatible_pattern(A, ag, uplo)
 
     # TODO: for some reason I cannot use ldiv! with a sparse QR
     strict_upper_nonzeros_A = M_factorization \ vec(B)

src/graph.jl

@@ -179,6 +179,7 @@ function build_edge_to_index(S::SparsityPatternCSC{T}) where {T}
     # edge_to_index gives an index for each edge
     edge_to_index = Vector{T}(undef, nnz(S))
     offsets = zeros(T, S.n)
+    nb_self_loops = 0
     counter = 0
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -193,10 +194,11 @@ function build_edge_to_index(S::SparsityPatternCSC{T}) where {T}
             elseif i == j
                 # this should never be used, make sure it errors
                 edge_to_index[k] = 0
+                nb_self_loops += 1
             end
         end
     end
-    return edge_to_index
+    return edge_to_index, nb_self_loops
 end
 
 ## Adjacency graph
@@ -227,16 +229,23 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 struct AdjacencyGraph{T<:Integer,augmented_graph}
     S::SparsityPatternCSC{T}
     edge_to_index::Vector{T}
+    nb_self_loops::Int
 end
 
 Base.eltype(::AdjacencyGraph{T}) where {T} = T
 
 function AdjacencyGraph(
     S::SparsityPatternCSC{T},
-    edge_to_index::Vector{T}=build_edge_to_index(S);
+    edge_to_index::Vector{T},
+    nb_self_loops::Int;
     augmented_graph::Bool=false,
 ) where {T}
-    return AdjacencyGraph{T,augmented_graph}(S, edge_to_index)
+    return AdjacencyGraph{T,augmented_graph}(S, edge_to_index, nb_self_loops)
+end
+
+function AdjacencyGraph(S::SparsityPatternCSC; augmented_graph::Bool=false)
+    edge_to_index, nb_self_loops = build_edge_to_index(S)
+    return AdjacencyGraph(S, edge_to_index, nb_self_loops; augmented_graph)
 end
 
 function AdjacencyGraph(A::SparseMatrixCSC; augmented_graph::Bool=false)
@@ -275,15 +284,7 @@ function degree(g::AdjacencyGraph{T,false}, v::Integer) where {T}
     return g.S.colptr[v + 1] - g.S.colptr[v] - has_selfloop
 end
 
-nb_edges(g::AdjacencyGraph{T,true}) where {T} = nnz(g.S) ÷ 2
-
-function nb_edges(g::AdjacencyGraph{T,false}) where {T}
-    ne = 0
-    for v in vertices(g)
-        ne += degree(g, v)
-    end
-    return ne ÷ 2
-end
+nb_edges(g::AdjacencyGraph) = (nnz(g.S) - g.nb_self_loops) ÷ 2
 
 maximum_degree(g::AdjacencyGraph) = maximum(Base.Fix1(degree, g), vertices(g))
 minimum_degree(g::AdjacencyGraph) = minimum(Base.Fix1(degree, g), vertices(g))

src/interface.jl

@@ -323,7 +323,7 @@ function _coloring(
     forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {R}
     A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; augmented_graph=true)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index, 0; augmented_graph=true)
     outputs_by_order = map(algo.orders) do order
         vertices_in_order = vertices(ag, order)
         _color, _star_set = star_coloring(
@@ -370,7 +370,7 @@ function _coloring(
     symmetric_pattern::Bool,
 ) where {R}
     A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; augmented_graph=true)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index, 0; augmented_graph=true)
     outputs_by_order = map(algo.orders) do order
         vertices_in_order = vertices(ag, order)
         _color, _tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)

src/matrices.jl

@@ -62,24 +62,59 @@ function respectful_similar(A::Union{Symmetric,Hermitian}, ::Type{T}) where {T}
 end
 
 """
-    same_pattern(A, B)
+    compatible_pattern(A::AbstractMatrix, bg::BipartiteGraph)
+    compatible_pattern(A::AbstractMatrix, ag::AdjacencyGraph, uplo::Symbol)
 
-Perform a partial equality check on the sparsity patterns of `A` and `B`:
+Perform a coarse compatibility check between the sparsity pattern of `A`
+and the reference sparsity pattern encoded in a graph structure.
 
-- if the return is `true`, they might have the same sparsity pattern but we're not sure
-- if the return is `false`, they definitely don't have the same sparsity pattern
+This function only checks necessary conditions for the two sparsity patterns to match.
+In particular, it may return `true` even if the patterns are not identical.
+
+When A is a `SparseMatrixCSC`, additional checks on the sparsity structure are performed.
+
+# Return value
+- `true`  : the sparsity patterns are potentially compatible
+- `false` : the sparsity patterns are definitely incompatible
 """
-same_pattern(A, B) = size(A) == size(B)
+compatible_pattern(A::AbstractMatrix, bg::BipartiteGraph) = size(A) == size(bg.S2)
+function compatible_pattern(A::SparseMatrixCSC, bg::BipartiteGraph)
+    size(A) == size(bg.S2) && nnz(A) == nnz(bg.S2)
+end
+
+function compatible_pattern(A::AbstractMatrix, ag::AdjacencyGraph, uplo::Symbol)
+    size(A) == size(ag.S)
+end
+function compatible_pattern(A::SparseMatrixCSC, ag::AdjacencyGraph, uplo::Symbol)
+    if uplo == :L || uplo == :U
+        return size(A) == size(ag.S) && nnz(A) == (nb_edges(ag) + ag.nb_self_loops)
+    else
+        return size(A) == size(ag.S) && nnz(A) == nnz(ag.S)
+    end
+end
 
-function same_pattern(
-    A::Union{SparseMatrixCSC,SparsityPatternCSC},
-    B::Union{SparseMatrixCSC,SparsityPatternCSC},
-)
-    return size(A) == size(B) && nnz(A) == nnz(B)
+function check_compatible_pattern(A::AbstractMatrix, bg::BipartiteGraph)
+    if !compatible_pattern(A, bg)
+        throw(DimensionMismatch("`A` and `bg.S2` must have the same sparsity pattern."))
+    end
 end
 
-function check_same_pattern(A, S)
-    if !same_pattern(A, S)
-        throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
+function check_compatible_pattern(A::AbstractMatrix, ag::AdjacencyGraph, uplo::Symbol)
+    if !compatible_pattern(A, ag, uplo)
+        if uplo == :L
+            throw(
+                DimensionMismatch(
+                    "`A` and `tril(ag.S)` must have the same sparsity pattern."
+                ),
+            )
+        elseif uplo == :U
+            throw(
+                DimensionMismatch(
+                    "`A` and `triu(ag.S)` must have the same sparsity pattern."
+                ),
+            )
+        else  # uplo == :F
+            throw(DimensionMismatch("`A` and `ag.S` must have the same sparsity pattern."))
+        end
     end
 end

src/result.jl

@@ -401,31 +401,31 @@ function TreeSetColoringResult(
     decompression_eltype::Type{R},
 ) where {T<:Integer,R}
     (; reverse_bfs_orders, tree_edge_indices, nt) = tree_set
-    (; S) = ag
+    (; S, nb_self_loops) = ag
     nvertices = length(color)
     group = group_by_color(T, color)
     rv = rowvals(S)
 
     # Vector for the decompression of the diagonal coefficients
-    diagonal_indices = T[]
-    diagonal_nzind = T[]
-    ndiag = 0
+    diagonal_indices = Vector{T}(undef, nb_self_loops)
+    diagonal_nzind = Vector{T}(undef, nb_self_loops)
 
     if !augmented_graph(ag)
+        l = 0
         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
+                    l += 1
+                    diagonal_indices[l] = i
+                    diagonal_nzind[l] = k
                 end
             end
         end
     end
 
     # Vectors for the decompression of the off-diagonal coefficients
-    nedges = (nnz(S) - ndiag) ÷ 2
+    nedges = nb_edges(ag)
     lower_triangle_offsets = Vector{T}(undef, nedges)
     upper_triangle_offsets = Vector{T}(undef, nedges)