Improve how we build A_and_Aᵀ (#156)

* Improve how we build A_and_Aᵀ

* Update src/interface.jl

* Add tests

* Doc

---------

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

Author: amontoison

docs/src/dev.md

@@ -16,6 +16,7 @@ SparseMatrixColorings.BipartiteGraph
 SparseMatrixColorings.vertices
 SparseMatrixColorings.neighbors
 transpose
+SparseMatrixColorings.bidirectional_pattern
 ```
 
 ## Low-level coloring

src/graph.jl

@@ -14,7 +14,7 @@ Copied from `SparseMatrixCSC`:
 - `colptr::Vector{Ti}`: column `j` is in `colptr[j]:(colptr[j+1]-1)`
 - `rowval::Vector{Ti}`: row indices of stored values
 """
-struct SparsityPatternCSC{Ti<:Integer}
+struct SparsityPatternCSC{Ti<:Integer} <: AbstractMatrix{Bool}
     m::Int
     n::Int
     colptr::Vector{Ti}
@@ -93,6 +93,76 @@ function Base.getindex(S::SparsityPatternCSC, i0::Integer, i1::Integer)
     return ((r1 > r2) || (rowvals(S)[r1] != i0)) ? false : true
 end
 
+"""
+    bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool)
+
+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(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
+
+function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern)
+    m, n = size(S)
+    p = m + n
+    nnzS = nnz(S)
+    rowval = Vector{Int}(undef, 2 * nnzS)
+    colptr = zeros(Int, p + 1)
+
+    # Update rowval and colptr for the block A
+    for i in 1:nnzS
+        rowval[i] = S.rowval[i] + n
+    end
+    for j in 1:n
+        colptr[j] = S.colptr[j]
+    end
+
+    # Update rowval and colptr for the block Aᵀ
+    if symmetric_pattern
+        # We use the sparsity pattern of A for Aᵀ
+        for i in 1:nnzS
+            rowval[nnzS + i] = S.rowval[i]
+        end
+        # m and n are identical because symmetric_pattern is true
+        for j in 1:m
+            colptr[n + j] = nnzS + S.colptr[j]
+        end
+        colptr[p + 1] = 2 * nnzS + 1
+    else
+        # We need to determine the sparsity pattern of Aᵀ
+        # We adapt the code of transpose(SparsityPatternCSC) in graph.jl
+        for k in 1:nnzS
+            i = S.rowval[k]
+            colptr[n + i] += 1
+        end
+
+        counter = 1
+        for col in (n + 1):p
+            nnz_col = colptr[col]
+            colptr[col] = counter
+            counter += nnz_col
+        end
+
+        for j in 1:n
+            for index in S.colptr[j]:(S.colptr[j + 1] - 1)
+                i = S.rowval[index]
+                pos = colptr[n + i]
+                rowval[nnzS + pos] = j
+                colptr[n + i] += 1
+            end
+        end
+
+        colptr[p + 1] = nnzS + counter
+        for col in p:-1:(n + 2)
+            colptr[col] = nnzS + colptr[col - 1]
+        end
+        colptr[n + 1] = nnzS + 1
+    end
+
+    # Create the SparsityPatternCSC of the augmented adjacency matrix
+    S_and_Sᵀ = SparsityPatternCSC{Int}(p, p, colptr, rowval)
+    return S_and_Sᵀ
+end
+
 ## Adjacency graph
 
 """

src/interface.jl

@@ -234,18 +234,9 @@ function coloring(
     decompression_eltype::Type{R}=Float64,
     symmetric_pattern::Bool=false,
 ) where {decompression,R}
-    m, n = size(A)
-    T = eltype(A)
-    Aᵀ = if symmetric_pattern || A isa Union{Symmetric,Hermitian}
-        A
-    else
-        transpose(A)
-    end  # TODO: fuse with next step?
-    A_and_Aᵀ = [
-        spzeros(T, n, n) SparseMatrixCSC(Aᵀ)
-        SparseMatrixCSC(A) spzeros(T, m, m)
-    ]  # TODO: slow
+    A_and_Aᵀ = bidirectional_pattern(A; symmetric_pattern)
     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)

test/graph.jl

@@ -4,6 +4,7 @@ using SparseMatrixColorings:
     SparsityPatternCSC,
     AdjacencyGraph,
     BipartiteGraph,
+    bidirectional_pattern,
     degree,
     degree_dist2,
     nb_vertices,
@@ -16,13 +17,33 @@ using Test
 @testset "SparsityPatternCSC" begin
     @testset "Transpose" begin
         @test all(1:1000) do _
-            A = sprand(rand(100:1000), rand(100:1000), 0.1)
+            m, n = rand(100:1000), rand(100:1000)
+            p = 0.05 * rand()
+            A = sprand(m, n, p)
             S = SparsityPatternCSC(A)
             Sᵀ = transpose(S)
             Sᵀ_true = SparsityPatternCSC(sparse(transpose(A)))
             Sᵀ.colptr == Sᵀ_true.colptr && Sᵀ.rowval == Sᵀ_true.rowval
         end
     end
+    @testset "Bidirectional" begin
+        @test all(1:1000) do _
+            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
+        end
+        @test all(1:1000) do _
+            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
+        end
+    end
     @testset "size" begin
         A = spzeros(10, 20)
         S = SparsityPatternCSC(A)