Update bidirectional_pattern

Author: amontoison

src/graph.jl

@@ -94,33 +94,44 @@ function Base.getindex(S::SparsityPatternCSC, i0::Integer, i1::Integer)
 end
 
 """
-    bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool)
+    bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool=false)
 
 Return a [`SparsityPatternCSC`](@ref) corresponding to the matrix `[0 Aᵀ; A 0]`, with a minimum of allocations.
 """
-bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool) =
+bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool=false) =
     bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
 
-function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern::Bool)
+function bidirectional_pattern(
+    S::SparsityPatternCSC{T}; symmetric_pattern::Bool=false
+) where {T<:Integer}
     m, n = size(S)
     p = m + n
     nnzS = nnz(S)
-    rowval = Vector{Int}(undef, 2 * nnzS)
-    colptr = zeros(Int, p + 1)
+    rowval = Vector{T}(undef, 2 * nnzS)
+    colptr = zeros(T, p + 1)
+    edge_to_index = Vector{T}(undef, 2 * nnzS)
 
     # Update rowval and colptr for the block A
     for i in 1:nnzS
         rowval[i] = S.rowval[i] + n
+        edge_to_index[i] = i
     end
     for j in 1:n
         colptr[j] = S.colptr[j]
     end
 
     # Update rowval and colptr for the block Aᵀ
     if symmetric_pattern
+        # Use colptr[n+1:n+m] to store the offsets during the update of edge_to_index
+        offsets = colptr
+
         # We use the sparsity pattern of A for Aᵀ
-        for i in 1:nnzS
-            rowval[nnzS + i] = S.rowval[i]
+        for k in 1:nnzS
+            r = S.rowval[k]
+            rowval[nnzS + k] = r
+            pos = S.colptr[r] + offsets[n + r]
+            edge_to_index[nnzS + pos] = edge_to_index[k]
+            offsets[n + r] += 1
         end
         # m and n are identical because symmetric_pattern is true
         for j in 1:m
@@ -147,6 +158,7 @@ function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern::Bool)
                 i = S.rowval[index]
                 pos = colptr[n + i]
                 rowval[nnzS + pos] = j
+                edge_to_index[nnzS + pos] = edge_to_index[index]
                 colptr[n + i] += 1
             end
         end
@@ -159,12 +171,10 @@ function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern::Bool)
     end
 
     # Create the SparsityPatternCSC of the augmented adjacency matrix
-    S_and_Sᵀ = SparsityPatternCSC{Int}(p, p, colptr, rowval)
-    return S_and_Sᵀ
+    S_and_Sᵀ = SparsityPatternCSC{T}(p, p, colptr, rowval)
+    return S_and_Sᵀ, edge_to_index
 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{T}) where {T<:Integer}
     # edge_to_index gives an index for each edge
     edge_to_index = Vector{T}(undef, nnz(S))
@@ -221,11 +231,27 @@ struct AdjacencyGraph{T,has_diagonal}
     vertex_buffer::Vector{T}
 end
 
-function AdjacencyGraph(S::SparsityPatternCSC{T}; has_diagonal::Bool=true) where {T}
-    edge_to_index, vertex_buffer = build_edge_to_index(S)
+function AdjacencyGraph(
+    S::SparsityPatternCSC{T},
+    edge_to_index::Vector{T},
+    vertex_buffer::Vector{T};
+    has_diagonal::Bool=true,
+) where {T}
     return AdjacencyGraph{T,has_diagonal}(S, edge_to_index, vertex_buffer)
 end
 
+function AdjacencyGraph(
+    S::SparsityPatternCSC{T}, edge_to_index::Vector{T}; has_diagonal::Bool=true
+) where {T}
+    vertex_buffer = Vector{Int}(undef, S.n)
+    return AdjacencyGraph(S, edge_to_index, vertex_buffer; has_diagonal)
+end
+
+function AdjacencyGraph(S::SparsityPatternCSC; has_diagonal::Bool=true)
+    edge_to_index, vertex_buffer = build_edge_to_index(S)
+    return AdjacencyGraph(S, edge_to_index, vertex_buffer; has_diagonal)
+end
+
 function AdjacencyGraph(A::SparseMatrixCSC; has_diagonal::Bool=true)
     return AdjacencyGraph(SparsityPatternCSC(A); has_diagonal)
 end
@@ -248,8 +274,9 @@ end
 
 function neighbors_with_edge_indices(g::AdjacencyGraph, v::Integer)
     S = pattern(g)
-    neighbors_v = view(rowvals(S), nzrange(S, v))
-    edges_indices_v = view(edge_indices(g), nzrange(S, v))
+    range_v = nzrange(S, v)
+    neighbors_v = view(rowvals(S), range_v)
+    edges_indices_v = view(edge_indices(g), range_v)
     return zip(neighbors_v, edges_indices_v)
 end
 
@@ -328,7 +355,7 @@ When `symmetric_pattern` is `true`, this construction is more efficient.
 
 > [_What Color Is Your Jacobian? SparsityPatternCSC Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-struct BipartiteGraph{T<:Integer}
+struct BipartiteGraph{T}
     S1::SparsityPatternCSC{T}
     S2::SparsityPatternCSC{T}
 end

src/interface.jl

@@ -265,7 +265,7 @@ function _coloring(
     decompression_eltype::Type,
     symmetric_pattern::Bool,
 )
-    ag = AdjacencyGraph(A)
+    ag = AdjacencyGraph(A; has_diagonal=true)
     color, star_set = star_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
         return StarSetColoringResult(A, ag, color, star_set)
@@ -282,7 +282,7 @@ function _coloring(
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
 ) where {R}
-    ag = AdjacencyGraph(A)
+    ag = AdjacencyGraph(A; has_diagonal=true)
     color, tree_set = acyclic_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
         return TreeSetColoringResult(A, ag, color, tree_set, R)
@@ -299,8 +299,8 @@ function _coloring(
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
 ) where {R}
-    A_and_Aᵀ = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ; has_diagonal=false)
+    A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false)
     color, star_set = star_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
         symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
@@ -319,8 +319,8 @@ function _coloring(
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
 ) where {R}
-    A_and_Aᵀ = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ; has_diagonal=false)
+    A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false)
     color, tree_set = acyclic_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
         symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R)

test/graph.jl

@@ -27,21 +27,27 @@ using Test
         end
     end
     @testset "Bidirectional" begin
-        @test all(1:1000) do _
+        @testset "symmetric_pattern = false" begin
             m, n = rand(100:1000), rand(100:1000)
             p = 0.05 * rand()
             A = sprand(Bool, m, n, p)
             A_and_Aᵀ = [spzeros(Bool, n, n) transpose(A); A spzeros(Bool, m, m)]
-            S_and_Sᵀ = bidirectional_pattern(A; symmetric_pattern=false)
-            S_and_Sᵀ.colptr == A_and_Aᵀ.colptr && S_and_Sᵀ.rowval == A_and_Aᵀ.rowval
+            S_and_Sᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern=false)
+            @test S_and_Sᵀ.colptr == A_and_Aᵀ.colptr
+            @test S_and_Sᵀ.rowval == A_and_Aᵀ.rowval
+            M = SparseMatrixCSC(m+n, m+n, S_and_Sᵀ.colptr, S_and_Sᵀ.rowval, edge_to_index)
+            @test issymmetric(M)
         end
-        @test all(1:1000) do _
+        @testset "symmetric_pattern = true" begin
             m = rand(100:1000)
             p = 0.05 * rand()
             A = sparse(Symmetric(sprand(Bool, m, m, p)))
             A_and_Aᵀ = [spzeros(Bool, m, m) transpose(A); A spzeros(Bool, m, m)]
-            S_and_Sᵀ = bidirectional_pattern(A; symmetric_pattern=true)
-            S_and_Sᵀ.colptr == A_and_Aᵀ.colptr && S_and_Sᵀ.rowval == A_and_Aᵀ.rowval
+            S_and_Sᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern=true)
+            @test S_and_Sᵀ.colptr == A_and_Aᵀ.colptr
+            @test S_and_Sᵀ.rowval == A_and_Aᵀ.rowval
+            M = SparseMatrixCSC(2*m, 2*m, S_and_Sᵀ.colptr, S_and_Sᵀ.rowval, edge_to_index)
+            @test issymmetric(M)
         end
     end
     @testset "size" begin