[bicoloring] Use Br and Bc directly for the decompression

Author: amontoison

docs/src/dev.md

@@ -26,6 +26,7 @@ SparseMatrixColorings.partial_distance2_coloring
 SparseMatrixColorings.star_coloring
 SparseMatrixColorings.acyclic_coloring
 SparseMatrixColorings.group_by_color
+SparseMatrixColorings.remap_colors
 SparseMatrixColorings.Forest
 SparseMatrixColorings.StarSet
 SparseMatrixColorings.TreeSet
@@ -39,8 +40,8 @@ SparseMatrixColorings.RowColoringResult
 SparseMatrixColorings.StarSetColoringResult
 SparseMatrixColorings.TreeSetColoringResult
 SparseMatrixColorings.LinearSystemColoringResult
-SparseMatrixColorings.BicoloringResult
-SparseMatrixColorings.remap_colors
+SparseMatrixColorings.StarSetBicoloringResult
+SparseMatrixColorings.TreeSetBicoloringResult
 ```
 
 ## Testing

src/decompression.jl

@@ -515,12 +515,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, tree_edge_indices, nt, buffer) = result
     (; S) = ag
     uplo == :F && check_same_pattern(A, S)
-    R = eltype(A)
     fill!(A, zero(R))
 
     if eltype(buffer) == R
@@ -714,61 +716,184 @@ function decompress!(
     return A
 end
 
-## BicoloringResult
-
-function _join_compressed!(result::BicoloringResult, Br::AbstractMatrix, Bc::AbstractMatrix)
-    #=
-    Say we have an original matrix `A` of size `(n, m)` and we build an augmented matrix `A_and_Aᵀ = [zeros(n, n) Aᵀ; A zeros(m, m)]`.
-    Its first `1:n` columns have the form `[zeros(n); A[:, j]]` and its following `n+1:n+m` columns have the form `[A[i, :]; zeros(m)]`.
-    The symmetric column coloring is performed on `A_and_Aᵀ` and the column-wise compression of `A_and_Aᵀ` should return a matrix `Br_and_Bc`.
-    But in reality, `Br_and_Bc` is computed as two partial compressions: the row-wise compression `Br` (corresponding to `Aᵀ`) and the columnwise compression `Bc` (corresponding to `A`).
-    Before symmetric decompression, we must reconstruct `Br_and_Bc` from `Br` and `Bc`, knowing that the symmetric colors (those making up `Br_and_Bc`) are present in either a row of `Br`, a column of `Bc`, or both.
-    Therefore, the column indices in `Br_and_Bc` don't necessarily match with the row indices in `Br` or the column indices in `Bc` since some colors may be missing in the partial compressions.
-    The columns of the top part of `Br_and_Bc` (rows `1:n`) are the rows of `Br`, interlaced with zero columns whenever the current color hasn't been used to color any row.
-    The columns of the bottom part of `Br_and_Bc` (rows `n+1:n+m`) are the columns of `Bc`, interlaced with zero columns whenever the current color hasn't been used to color any column.
-    We use the vectors `symmetric_to_row` and `symmetric_to_column` to map from symmetric colors to row and column colors.
-    =#
-    (; A, symmetric_to_column, symmetric_to_row) = result
-    m, n = size(A)
-    R = Base.promote_eltype(Br, Bc)
-    if eltype(result.Br_and_Bc) == R
-        Br_and_Bc = result.Br_and_Bc
-    else
-        Br_and_Bc = similar(result.Br_and_Bc, R)
-    end
-    fill!(Br_and_Bc, zero(R))
-    for c in axes(Br_and_Bc, 2)
-        if symmetric_to_row[c] > 0  # some rows were colored with the symmetric color c
-            copyto!(view(Br_and_Bc, 1:n, c), view(Br, symmetric_to_row[c], :))
-        end
-        if symmetric_to_column[c] > 0  # some columns were colored with the symmetric color c
-            copyto!(
-                view(Br_and_Bc, (n + 1):(n + m), c), view(Bc, :, symmetric_to_column[c])
-            )
+## StarSetBicoloringResult
+
+function decompress!(
+    A::AbstractMatrix,
+    Br::AbstractMatrix,
+    Bc::AbstractMatrix,
+    result::StarSetBicoloringResult,
+)
+    (; S, A_indices, compressed_indices, pos_Bc) = result
+    fill!(A, zero(eltype(A)))
+
+    ind_Bc = 1
+    ind_Br = nnz(S)
+    rvS = rowvals(S)
+    for j in axes(S, 2)
+        for k in nzrange(S, j)
+            i = rvS[k]
+            index_Bc = compressed_indices[ind_Bc]
+            if A_indices[ind_Bc] == k
+                A[i, j] = Bc[index_Bc]
+                if ind_Bc < pos_Bc
+                    ind_Bc += 1
+                end
+            else
+                index_Br = compressed_indices[ind_Br]
+                A[i, j] = Br[index_Br]
+                ind_Br -= 1
+            end
         end
     end
-    return Br_and_Bc
+    return A
 end
 
 function decompress!(
-    A::AbstractMatrix, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult
+    A::SparseMatrixCSC,
+    Br::AbstractMatrix,
+    Bc::AbstractMatrix,
+    result::StarSetBicoloringResult,
 )
+    (; A_indices, compressed_indices, pos_Bc) = result
+    nzA = nonzeros(A)
+    for k in 1:pos_Bc
+        nzA[A_indices[k]] = Bc[compressed_indices[k]]
+    end
+    for k in (pos_Bc + 1):length(nzA)
+        nzA[A_indices[k]] = Br[compressed_indices[k]]
+    end
+    return A
+end
+
+## TreeSetBicoloringResult
+
+function decompress!(
+    A::AbstractMatrix{R},
+    Br::AbstractMatrix{R},
+    Bc::AbstractMatrix{R},
+    result::TreeSetBicoloringResult,
+) where {R<:Real}
+    (;
+        symmetric_color,
+        symmetric_to_row,
+        symmetric_to_column,
+        reverse_bfs_orders,
+        tree_edge_indices,
+        nt,
+        buffer,
+    ) = result
+
     m, n = size(A)
-    Br_and_Bc = _join_compressed!(result, Br, Bc)
-    A_and_Aᵀ = decompress(Br_and_Bc, result.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 1:nt
+        # Positions of the first and last edges of the tree
+        first = tree_edge_indices[k]
+        last = tree_edge_indices[k + 1] - 1
+
+        # Reset the buffer to zero for all vertices in the tree (except the root)
+        for pos in first:last
+            (vertex, _) = reverse_bfs_orders[pos]
+            buffer_right_type[vertex] = zero(R)
+        end
+        # Reset the buffer to zero for the root vertex
+        (_, root) = reverse_bfs_orders[last]
+        buffer_right_type[root] = zero(R)
+
+        for pos in first:last
+            (i, j) = reverse_bfs_orders[pos]
+            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, result::BicoloringResult
-)
-    (; large_colptr, large_rowval, symmetric_result) = result
+    A::SparseMatrixCSC{R},
+    Br::AbstractMatrix{R},
+    Bc::AbstractMatrix{R},
+    result::TreeSetBicoloringResult,
+) where {R<:Real}
+    (;
+        symmetric_color,
+        symmetric_to_column,
+        symmetric_to_row,
+        reverse_bfs_orders,
+        tree_edge_indices,
+        nt,
+        A_indices,
+        buffer,
+    ) = result
+
     m, n = size(A)
-    Br_and_Bc = _join_compressed!(result, Br, Bc)
-    # 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)
+    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
+
+    # Recover the off-diagonal coefficients of A
+    for k in 1:nt
+        # Positions of the first and last edges of the tree
+        first = tree_edge_indices[k]
+        last = tree_edge_indices[k + 1] - 1
+
+        # Reset the buffer to zero for all vertices in the tree (except the root)
+        for pos in first:last
+            (vertex, _) = reverse_bfs_orders[pos]
+            buffer_right_type[vertex] = zero(R)
+        end
+        # Reset the buffer to zero for the root vertex
+        (_, root) = reverse_bfs_orders[last]
+        buffer_right_type[root] = zero(R)
+
+        for pos in first:last
+            (i, j) = reverse_bfs_orders[pos]
+            cj = symmetric_color[j]
+            counter += 1
+
+            #! 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/graph.jl

@@ -100,7 +100,8 @@ end
 Return a [`SparsityPatternCSC`](@ref) corresponding to the matrix `[0 Aᵀ; A 0]`, with a minimum of allocations.
 """
 function bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool)
-    bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
+    S = SparsityPatternCSC(SparseMatrixCSC(A))
+    bidirectional_pattern(S; symmetric_pattern)
 end
 
 function bidirectional_pattern(S::SparsityPatternCSC{T}; symmetric_pattern::Bool) where {T}
@@ -172,7 +173,7 @@ function bidirectional_pattern(S::SparsityPatternCSC{T}; symmetric_pattern::Bool
 
     # Create the SparsityPatternCSC of the augmented adjacency matrix
     S_and_Sᵀ = SparsityPatternCSC{T}(p, p, colptr, rowval)
-    return S_and_Sᵀ, edge_to_index
+    return S, S_and_Sᵀ, edge_to_index
 end
 
 function build_edge_to_index(S::SparsityPatternCSC{T}) where {T}

src/interface.jl

@@ -282,7 +282,7 @@ function _coloring(
     algo::GreedyColoringAlgorithm{:substitution},
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
-) where {R}
+) where {R<:Real}
     ag = AdjacencyGraph(A; has_diagonal=true)
     vertices_in_order = vertices(ag, algo.order)
     color, tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
@@ -298,19 +298,18 @@ function _coloring(
     A::AbstractMatrix,
     ::ColoringProblem{:nonsymmetric,:bidirectional},
     algo::GreedyColoringAlgorithm{:direct},
-    decompression_eltype::Type{R},
+    decompression_eltype::Type,
     symmetric_pattern::Bool,
 ) where {R}
-    A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false)
+    S, S_and_Sᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
+    ag = AdjacencyGraph(S_and_Sᵀ, edge_to_index; has_diagonal=false)
     vertices_in_order = vertices(ag, algo.order)
-    color, star_set = star_coloring(ag, vertices_in_order, algo.postprocessing)
+    symmetric_color, star_set = star_coloring(ag, vertices_in_order, algo.postprocessing)
     if speed_setting isa WithResult
-        symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
-        return BicoloringResult(A, ag, symmetric_result, R)
+        return StarSetBicoloringResult(A, S, ag, symmetric_color, star_set)
     else
         row_color, column_color, _ = remap_colors(
-            eltype(ag), color, maximum(color), size(A)...
+            eltype(ag), symmetric_color, maximum(symmetric_color), size(A)...
         )
         return row_color, column_color
     end
@@ -324,16 +323,15 @@ function _coloring(
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
 ) where {R}
-    A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false)
+    S, S_and_Sᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
+    ag = AdjacencyGraph(S_and_Sᵀ, edge_to_index; has_diagonal=false)
     vertices_in_order = vertices(ag, algo.order)
-    color, tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
+    symmetric_color, tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
     if speed_setting isa WithResult
-        symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R)
-        return BicoloringResult(A, ag, symmetric_result, R)
+        return TreeSetBicoloringResult(A, S, ag, symmetric_color, tree_set, R)
     else
         row_color, column_color, _ = remap_colors(
-            eltype(ag), color, maximum(color), size(A)...
+            eltype(ag), symmetric_color, maximum(symmetric_color), size(A)...
         )
         return row_color, column_color
     end