Save allocations for bipartite graph with symmetric patterns (#102)

Author: gdalle

src/graph.jl

@@ -79,7 +79,7 @@ end
 ## Construct from matrices
 
 """
-    adjacency_graph(H::AbstractMatrix)
+    adjacency_graph(A::SparseMatrixCSC)
 
 Return a [`Graph`](@ref) representing the nonzeros of a symmetric matrix (typically a Hessian matrix).
 
@@ -92,10 +92,10 @@ The adjacency graph of a symmetrix matric `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-adjacency_graph(H::SparseMatrixCSC) = Graph(H - Diagonal(H))
+adjacency_graph(A::SparseMatrixCSC) = Graph(A - Diagonal(A))
 
 """
-    bipartite_graph(J::AbstractMatrix)
+    bipartite_graph(A::SparseMatrixCSC; symmetric_pattern::Bool)
 
 Return a [`BipartiteGraph`](@ref) representing the nonzeros of a non-symmetric matrix (typically a Jacobian matrix).
 
@@ -105,12 +105,19 @@ The bipartite graph of a matrix `A ∈ ℝ^{m × n}` is `Gb(A) = (V₁, V₂, E)
 - `V₂ = 1:n` is the set of columns `j`
 - `(i, j) ∈ E` whenever `A[i, j] ≠ 0`
 
+When `symmetric_pattern` is `true`, this construction is more efficient.
+
 # References
 
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-function bipartite_graph(J::SparseMatrixCSC)
-    g1 = Graph(sparse(transpose(J)))  # rows to columns
-    g2 = Graph(J)  # columns to rows
+function bipartite_graph(A::SparseMatrixCSC; symmetric_pattern::Bool=false)
+    g2 = Graph(A)  # columns to rows
+    if symmetric_pattern
+        checksquare(A)  # proxy for checking full symmetry
+        g1 = g2
+    else
+        g1 = Graph(sparse(transpose(A)))  # rows to columns
+    end
     return BipartiteGraph(g1, g2)
 end

src/interface.jl

@@ -117,14 +117,18 @@ end
         S::AbstractMatrix,
         problem::ColoringProblem,
         algo::GreedyColoringAlgorithm;
-        [decompression_eltype=Float64]
+        [decompression_eltype=Float64, symmetric_pattern=false]
     )
 
 Solve a [`ColoringProblem`](@ref) on the matrix `S` with a [`GreedyColoringAlgorithm`](@ref) and return an [`AbstractColoringResult`](@ref).
 
 The result can be used to [`compress`](@ref) and [`decompress`](@ref) a matrix `A` with the same sparsity pattern as `S`.
 If `eltype(A) == decompression_eltype`, decompression might be faster.
 
+For a `:nonsymmetric` problem (and only then), setting `symmetric_pattern=true` indicates that the pattern of nonzeros is symmetric.
+This condition is weaker than the symmetry of actual values, so it can happen for some Jacobians.
+Specifying it allows faster construction of the bipartite graph.
+
 # Example
 
 ```jldoctest
@@ -173,10 +177,13 @@ function coloring(
     A::AbstractMatrix,
     ::ColoringProblem{:nonsymmetric,:column},
     algo::GreedyColoringAlgorithm;
-    decompression_eltype=Float64,
+    decompression_eltype::Type=Float64,
+    symmetric_pattern::Bool=false,
 )
     S = sparse(A)
-    bg = bipartite_graph(S)
+    bg = bipartite_graph(
+        S; symmetric_pattern=symmetric_pattern || A isa Union{Symmetric,Hermitian}
+    )
     color = partial_distance2_coloring(bg, Val(2), algo.order)
     return ColumnColoringResult(S, color)
 end
@@ -185,10 +192,13 @@ function coloring(
     A::AbstractMatrix,
     ::ColoringProblem{:nonsymmetric,:row},
     algo::GreedyColoringAlgorithm;
-    decompression_eltype=Float64,
+    decompression_eltype::Type=Float64,
+    symmetric_pattern::Bool=false,
 )
     S = sparse(A)
-    bg = bipartite_graph(S)
+    bg = bipartite_graph(
+        S; symmetric_pattern=symmetric_pattern || A isa Union{Symmetric,Hermitian}
+    )
     color = partial_distance2_coloring(bg, Val(1), algo.order)
     return RowColoringResult(S, color)
 end
@@ -197,7 +207,7 @@ function coloring(
     A::AbstractMatrix,
     ::ColoringProblem{:symmetric,:column},
     algo::GreedyColoringAlgorithm{:direct};
-    decompression_eltype=Float64,
+    decompression_eltype::Type=Float64,
 )
     S = sparse(A)
     ag = adjacency_graph(S)
@@ -209,7 +219,7 @@ function coloring(
     A::AbstractMatrix,
     ::ColoringProblem{:symmetric,:column},
     algo::GreedyColoringAlgorithm{:substitution};
-    decompression_eltype=Float64,
+    decompression_eltype::Type=Float64,
 )
     S = sparse(A)
     ag = adjacency_graph(S)
@@ -221,14 +231,14 @@ end
 
 function ADTypes.column_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
     S = sparse(A)
-    bg = bipartite_graph(S)
+    bg = bipartite_graph(S; symmetric_pattern=A isa Union{Symmetric,Hermitian})
     color = partial_distance2_coloring(bg, Val(2), algo.order)
     return color
 end
 
 function ADTypes.row_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
     S = sparse(A)
-    bg = bipartite_graph(S)
+    bg = bipartite_graph(S; symmetric_pattern=A isa Union{Symmetric,Hermitian})
     color = partial_distance2_coloring(bg, Val(1), algo.order)
     return color
 end

test/graph.jl

@@ -28,7 +28,8 @@ end;
         0 0 0 1 1 1 1 0
     ])
 
-    bg = bipartite_graph(A)
+    bg = bipartite_graph(A; symmetric_pattern=false)
+    @test_throws DimensionMismatch bipartite_graph(A; symmetric_pattern=true)
     @test length(bg, Val(1)) == 4
     @test length(bg, Val(2)) == 8
     # neighbors of rows
@@ -45,6 +46,21 @@ end;
     @test neighbors(bg, Val(2), 6) == [1, 3, 4]
     @test neighbors(bg, Val(2), 7) == [1, 2, 4]
     @test neighbors(bg, Val(2), 8) == [1, 2, 3]
+
+    A = sparse([
+        1 0 1 1
+        0 1 0 1
+        1 0 1 0
+        1 1 0 1
+    ])
+    bg = bipartite_graph(A; symmetric_pattern=true)
+    @test length(bg, Val(1)) == 4
+    @test length(bg, Val(2)) == 4
+    # neighbors of rows and columns
+    @test neighbors(bg, Val(1), 1) == neighbors(bg, Val(2), 1) == [1, 3, 4]
+    @test neighbors(bg, Val(1), 2) == neighbors(bg, Val(2), 2) == [2, 4]
+    @test neighbors(bg, Val(1), 3) == neighbors(bg, Val(2), 3) == [1, 3]
+    @test neighbors(bg, Val(1), 4) == neighbors(bg, Val(2), 4) == [1, 2, 4]
 end;
 
 ## Adjacency graph (fig 3.1 of "What color is your Jacobian?")

test/random.jl

@@ -33,6 +33,13 @@ symmetric_params = vcat(
         @test directly_recoverable_columns(A0, color0)
         test_coloring_decompression(A0, problem, algo; color0)
     end
+    @testset "$((; n, p))" for (n, p) in symmetric_params
+        A0 = sparse(Symmetric(sprand(rng, n, n, p)))
+        color0 = column_coloring(A0, algo)
+        @test structurally_orthogonal_columns(A0, color0)
+        @test directly_recoverable_columns(A0, color0)
+        test_coloring_decompression(A0, problem, algo; color0)
+    end
 end;
 
 @testset "Row coloring & decompression" begin
@@ -45,6 +52,13 @@ end;
         @test directly_recoverable_columns(transpose(A0), color0)
         test_coloring_decompression(A0, problem, algo; color0)
     end
+    @testset "$((; n, p))" for (n, p) in symmetric_params
+        A0 = sparse(Symmetric(sprand(rng, n, n, p)))
+        color0 = row_coloring(A0, algo)
+        @test structurally_orthogonal_columns(transpose(A0), color0)
+        @test directly_recoverable_columns(transpose(A0), color0)
+        test_coloring_decompression(A0, problem, algo; color0)
+    end
 end;
 
 @testset "Symmetric coloring & direct decompression" begin

test/utils.jl

@@ -12,7 +12,15 @@ function test_coloring_decompression(
 ) where {structure,partition,decompression}
     color_vec = Vector{Int}[]
     @testset "$(typeof(A))" for A in matrix_versions(A0)
-        result = coloring(A, problem, algo; decompression_eltype=Float64)
+        yield()
+
+        if structure == :nonsymmetric && issymmetric(A)
+            result = coloring(
+                A, problem, algo; decompression_eltype=Float64, symmetric_pattern=true
+            )
+        else
+            result = coloring(A, problem, algo; decompression_eltype=Float64)
+        end
         color = if partition == :column
             column_colors(result)
         elseif partition == :row