Decompression for TreeSetBicoloringResult

Author: amontoison

docs/src/dev.md

@@ -27,6 +27,7 @@ SparseMatrixColorings.symmetric_coefficient
 SparseMatrixColorings.star_coloring
 SparseMatrixColorings.acyclic_coloring
 SparseMatrixColorings.group_by_color
+SparseMatrixColorings.remap_colors
 SparseMatrixColorings.StarSet
 SparseMatrixColorings.TreeSet
 ```
@@ -43,13 +44,6 @@ SparseMatrixColorings.StarSetBicoloringResult
 SparseMatrixColorings.TreeSetBicoloringResult
 ```
 
-## Decompression
-
-```@docs
-SparseMatrixColorings.JoinCompressed
-SparseMatrixColorings.remap_colors
-```
-
 ## Testing
 
 ```@docs

src/decompression.jl

@@ -531,12 +531,14 @@ end
 ## TreeSetColoringResult
 
 function decompress!(
-    A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
-)
+    A::AbstractMatrix{R},
+    B::AbstractMatrix{R},
+    result::TreeSetColoringResult,
+    uplo::Symbol=:F,
+) where {R<:Real}
     (; ag, color, reverse_bfs_orders, buffer) = result
     (; S) = ag
     uplo == :F && check_same_pattern(A, S)
-    R = eltype(A)
     fill!(A, zero(R))
 
     if eltype(buffer) == R
@@ -715,49 +717,6 @@ function decompress!(
     return A
 end
 
-## BicoloringResult
-
-"""
-    JoinCompressed{R<:Real}
-
-For a bicoloring of an original matrix `A` of size `(m, n)`, we build an augmented matrix `A_and_Aᵀ = [0 Aᵀ; A 0]`.
-Symmetric column coloring is then performed on `A_and_Aᵀ`, and the column-wise compression of `A_and_Aᵀ` produces a matrix `Br_and_Bc`.
-In the case of bicoloring, `Br_and_Bc` is computed as two partial compressions: the row-wise compression `Br` (corresponding to `Aᵀ`) and the column-wise compression `Bc` (corresponding to `A`).
-
-For the symmetric decompression, we lazily reconstruct `Br_and_Bc` from `Br` and `Bc`, knowing that the symmetric colors (those making up `Br_and_Bc`) may appear in either a row of `Br`, a column of `Bc`, or both.
-The columns of the top part of `Br_and_Bc` (rows between `1` and `n`) are taken from the rows of `Br`, interleaved with zero columns whenever the current color has not been used to color any row.
-The columns of the bottom part of `Br_and_Bc` (rows between `n+1` and `n+m`) are taken from the columns of `Bc`, interleaved with zero columns whenever the current color has not been used to color any column.
-
-We use the vectors `symmetric_to_row` and `symmetric_to_column` to map colors obtained during star or acyclic coloring to row colors in `Br` and column colors in `Bc`.
-
-# Fields
-
-- `m::Int`: number of rows in `A` and `Bc`
-- `n::Int`: number of columns in `A` and `Br`
-- `Br::AbstractMatrix{R}`: row-wise compressed matrix
-- `Bc::AbstractMatrix{R}`: column-wise compressed matrix
-- `symmetric_to_row::Vector{Int}`: vector mapping symmetric colors to row indices in `Br`
-- `symmetric_to_column::Vector{Int}`: vector mapping symmetric colors to column indices in `Bc`
-"""
-struct JoinCompressed{R<:Real,M1<:AbstractMatrix{R},M2<:AbstractMatrix{R}} <:
-       AbstractMatrix{R}
-    m::Int
-    n::Int
-    Br::M1
-    Bc::M2
-    symmetric_to_row::Vector{Int}
-    symmetric_to_column::Vector{Int}
-end
-
-function Base.getindex(B::JoinCompressed, i::Int, j::Int)
-    (; n, Br, Bc, symmetric_to_row, symmetric_to_column) = B
-    if i ≤ n
-        return Br[symmetric_to_row[j], i]
-    else
-        return Bc[i - n, symmetric_to_column[j]]
-    end
-end
-
 ## StarSetBicoloringResult
 
 function decompress!(
@@ -807,33 +766,118 @@ end
 ## TreeSetBicoloringResult
 
 function decompress!(
-    A::AbstractMatrix,
-    Br::AbstractMatrix,
-    Bc::AbstractMatrix,
+    A::AbstractMatrix{R},
+    Br::AbstractMatrix{R},
+    Bc::AbstractMatrix{R},
     result::TreeSetBicoloringResult,
-)
-    (; symmetric_to_row, symmetric_to_column, symmetric_result) = result
+) where {R<:Real}
+    (;
+        ag,
+        symmetric_color,
+        symmetric_to_row,
+        symmetric_to_column,
+        reverse_bfs_orders,
+        buffer,
+    ) = result
+    (; S) = ag
+
     m, n = size(A)
-    Br_and_Bc = JoinCompressed(m, n, Br, Bc, symmetric_to_row, symmetric_to_column)
-    A_and_Aᵀ = decompress(Br_and_Bc, symmetric_result)
-    copyto!(A, A_and_Aᵀ[(n + 1):(n + m), 1:n])  # original matrix in bottom left corner
+    fill!(A, zero(R))
+
+    if eltype(buffer) == R
+        buffer_right_type = buffer
+    else
+        buffer_right_type = similar(buffer, R)
+    end
+
+    for k in eachindex(reverse_bfs_orders)
+        # Reset the buffer to zero for all vertices in a tree (except the root)
+        for (vertex, _) in reverse_bfs_orders[k]
+            buffer_right_type[vertex] = zero(R)
+        end
+        # Reset the buffer to zero for the root vertex
+        (_, root) = reverse_bfs_orders[k][end]
+        buffer_right_type[root] = zero(R)
+
+        for (i, j) in reverse_bfs_orders[k]
+            cj = symmetric_color[j]
+            if in_triangle(i, j, :L)
+                val = Bc[i - n, symmetric_to_column[cj]] - buffer_right_type[i]
+                buffer_right_type[j] = buffer_right_type[j] + val
+                A[i - n, j] = val
+            else
+                val = Br[symmetric_to_row[cj], i] - buffer_right_type[i]
+                buffer_right_type[j] = buffer_right_type[j] + val
+                A[j - n, i] = val
+            end
+        end
+    end
     return A
 end
 
 function decompress!(
-    A::SparseMatrixCSC,
-    Br::AbstractMatrix,
-    Bc::AbstractMatrix,
+    A::SparseMatrixCSC{R},
+    Br::AbstractMatrix{R},
+    Bc::AbstractMatrix{R},
     result::TreeSetBicoloringResult,
-)
+) where {R<:Real}
     (;
-        symmetric_to_row, symmetric_to_column, symmetric_result, large_colptr, large_rowval
+        ag,
+        symmetric_color,
+        symmetric_to_column,
+        symmetric_to_row,
+        reverse_bfs_orders,
+        A_indices,
+        buffer,
     ) = result
+    (; S) = ag
+
     m, n = size(A)
-    Br_and_Bc = JoinCompressed(m, n, Br, Bc, symmetric_to_row, symmetric_to_column)
-    # pretend A is larger
-    A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, A.nzval)
-    # decompress lower triangle only
-    decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L)
+    A_colptr = A.colptr
+    nzA = nonzeros(A)
+
+    if eltype(buffer) == R
+        buffer_right_type = buffer
+    else
+        buffer_right_type = similar(buffer, R)
+    end
+
+    # Recover the coefficients of A
+    counter = 0
+    for k in eachindex(reverse_bfs_orders)
+        # Reset the buffer to zero for all vertices in a tree (except the root)
+        for (vertex, _) in reverse_bfs_orders[k]
+            buffer_right_type[vertex] = zero(R)
+        end
+        # Reset the buffer to zero for the root vertex
+        (_, root) = reverse_bfs_orders[k][end]
+        buffer_right_type[root] = zero(R)
+
+        for (i, j) in reverse_bfs_orders[k]
+            counter += 1
+            cj = symmetric_color[j]
+
+            #! format: off
+            # A[i,j] is in the lower triangular part of A
+            if in_triangle(i, j, :L)
+                val = Bc[i - n, symmetric_to_column[cj]] - buffer_right_type[i]
+                buffer_right_type[j] = buffer_right_type[j] + val
+
+                # A[i,j] is stored at index_ij = A_indices[counter] in A.nzval
+                nzind = A_indices[counter]
+                nzA[nzind] = val
+
+            # A[i,j] is in the upper triangular part of A
+            else
+                val = Br[symmetric_to_row[cj], i] - buffer_right_type[i]
+                buffer_right_type[j] = buffer_right_type[j] + val
+
+                # A[j,i] is stored at index_ji = A_indices[counter] in A.nzval
+                nzind = A_indices[counter]
+                nzA[nzind] = val
+            end
+            #! format: on
+        end
+    end
     return A
 end

src/interface.jl

@@ -328,8 +328,7 @@ function _coloring(
         ag, algo.order; postprocessing=algo.postprocessing
     )
     if speed_setting isa WithResult
-        symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, symmetric_color, tree_set, R)
-        return TreeSetBicoloringResult(A, ag, symmetric_result)
+        return TreeSetBicoloringResult(A, ag, symmetric_color, tree_set, R)
     else
         row_color, column_color, _ = remap_colors(
             symmetric_color, maximum(symmetric_color), size(A)...

src/result.jl

@@ -292,7 +292,7 @@ function TreeSetColoringResult(
     color::Vector{Int},
     tree_set::TreeSet,
     decompression_eltype::Type{R},
-) where {R}
+) where {R<:Real}
     (; reverse_bfs_orders) = tree_set
     (; S) = ag
     nvertices = length(color)
@@ -596,7 +596,6 @@ function StarSetBicoloringResult(
             end
         end
     end
-    @assert pos_Bc - pos_Br == 1
 
     return StarSetBicoloringResult(
         A,
@@ -628,16 +627,14 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct TreeSetBicoloringResult{
-    M<:AbstractMatrix,
-    G<:AdjacencyGraph,
-    V,
-    SR<:AbstractColoringResult{:symmetric,:column,:substitution},
-} <: AbstractColoringResult{:nonsymmetric,:bidirectional,:substitution}
+struct TreeSetBicoloringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V,R} <:
+       AbstractColoringResult{:nonsymmetric,:bidirectional,:substitution}
     "matrix that was colored"
     A::M
     "augmented adjacency graph that was used for bicoloring"
     ag::G
+    "one integer color for each axis of the augmented matrix"
+    symmetric_color::Vector{Int}
     "one integer color for each column"
     column_color::Vector{Int}
     "one integer color for each row"
@@ -646,45 +643,80 @@ struct TreeSetBicoloringResult{
     column_group::V
     "color groups for rows"
     row_group::V
-    "result for the coloring of the symmetric 2 x 2 block matrix"
-    symmetric_result::SR
     "maps symmetric colors to column colors"
     symmetric_to_column::Vector{Int}
     "maps symmetric colors to row colors"
     symmetric_to_row::Vector{Int}
-    "CSC storage of `A_and_noAᵀ - `colptr`"
-    large_colptr::Vector{Int}
-    "CSC storage of `A_and_noAᵀ - `rowval`"
-    large_rowval::Vector{Int}
+    "indices of the nonzeros in A that can recovered with Br and Bc"
+    A_indices::Vector{Int}
+    "reverse BFS orders of the trees"
+    reverse_bfs_orders::Vector{Vector{Tuple{Int,Int}}}
+    "buffer needed during acyclic decompression"
+    buffer::Vector{R}
 end
 
 function TreeSetBicoloringResult(
     A::AbstractMatrix,
     ag::AdjacencyGraph,
-    symmetric_result::AbstractColoringResult{:symmetric,:column,:substitution},
-)
+    symmetric_color::Vector{Int},
+    tree_set::TreeSet,
+    decompression_eltype::Type{R},
+) where {R<:Real}
+    (; reverse_bfs_orders) = tree_set
+
     m, n = size(A)
-    symmetric_color = column_colors(symmetric_result)
     num_sym_colors = maximum(symmetric_color)
     row_color, column_color, symmetric_to_row, symmetric_to_column = remap_colors(
         symmetric_color, num_sym_colors, m, n
     )
     column_group = group_by_color(column_color)
     row_group = group_by_color(row_color)
-    large_colptr = copy(ag.S.colptr)
-    large_colptr[(n + 2):end] .= large_colptr[n + 1]  # last few columns are empty
-    large_rowval = ag.S.rowval[1:(end ÷ 2)]  # forget the second half of nonzeros
+
+    S = ag.S
+    rv = rowvals(S)
+    nnzA = nnz(S) ÷ 2
+    A_indices = Vector{Int}(undef, nnzA)
+
+    index = 0
+    for k in eachindex(reverse_bfs_orders)
+        for (i, j) in reverse_bfs_orders[k]
+            index += 1
+
+            #! format: off
+            # S[i,j] is in the lower triangular part of S
+            if in_triangle(i, j, :L)
+                # S[i,j] is stored at index_ij = (S.colptr[j] + offset) in S.nzval
+                col_j = view(rv, nzrange(S, j))
+                offset = searchsortedfirst(col_j, i)::Int - 1
+                A_indices[index] = S.colptr[j] + offset
+
+            # S[i,j] is in the upper triangular part of S
+            else
+                # S[j,i] is stored at index_ji = (S.colptr[i] + offset) in S.nzval
+                col_i = view(rv, nzrange(S, i))
+                offset = searchsortedfirst(col_i, j)::Int - 1
+                A_indices[index] = S.colptr[i] + offset
+            end
+            #! format: on
+        end
+    end
+
+    # buffer holds the sum of edge values for subtrees in a tree.
+    # For each vertex i, buffer[i] is the sum of edge values in the subtree rooted at i.
+    buffer = Vector{R}(undef, n + m)
+
     return TreeSetBicoloringResult(
         A,
         ag,
+        symmetric_color,
         column_color,
         row_color,
         column_group,
         row_group,
-        symmetric_result,
         symmetric_to_column,
         symmetric_to_row,
-        large_colptr,
-        large_rowval,
+        A_indices,
+        reverse_bfs_orders,
+        buffer,
     )
 end

test/allocations.jl

@@ -35,11 +35,7 @@ function test_noallocs_sparse_decompression(
             bench1_full = @be similar(A) decompress!(_, Br, Bc, result) evals = 1
             bench2_full = @be similar(Matrix(A)) decompress!(_, Br, Bc, result) evals = 1
             @test minimum(bench1_full).allocs == 0
-            if decompression == :direct
-                @test minimum(bench2_full).allocs == 0
-            else
-                @test_broken minimum(bench2_full).allocs == 0
-            end
+            @test minimum(bench2_full).allocs == 0
         end
     else
         B = compress(A, result)