Add StarSetBicoloringResult and TreeSetBicoloringResult

Author: amontoison

docs/src/dev.md

@@ -39,7 +39,8 @@ SparseMatrixColorings.RowColoringResult
 SparseMatrixColorings.StarSetColoringResult
 SparseMatrixColorings.TreeSetColoringResult
 SparseMatrixColorings.LinearSystemColoringResult
-SparseMatrixColorings.BicoloringResult
+SparseMatrixColorings.StarSetBicoloringResult
+SparseMatrixColorings.TreeSetBicoloringResult
 ```
 
 ## Decompression

src/decompression.jl

@@ -528,41 +528,6 @@ function decompress!(
     return A
 end
 
-function decompress!(
-    A::SparseMatrixCSC,
-    Br::AbstractMatrix,
-    Bc::AbstractMatrix,
-    symmetric_to_row::Vector{Int},
-    symmetric_to_column::Vector{Int},
-    result::StarSetColoringResult,
-)
-    (; ag, color, compressed_indices) = result
-    (; S) = ag
-    n = size(Br, 2)
-    m = size(Bc, 1)
-    dim = m + n
-    nzA = nonzeros(A)
-    rvS = rowvals(S)
-    l = 0  # assume A has the same pattern as the triangle
-    for j in axes(S, 2)
-        for k in nzrange(S, j)
-            i = rvS[k]
-            if in_triangle(i, j, :L)
-                l += 1
-                j2, i2 = divrem(compressed_indices[k] - 1, dim)
-                j2 += 1
-                i2 += 1
-                if i2 ≤ n
-                    nzA[l] = Br[symmetric_to_row[j2], i2]
-                else
-                    nzA[l] = Bc[i2 - n, symmetric_to_column[j2]]
-                end
-            end
-        end
-    end
-    return A
-end
-
 ## TreeSetColoringResult
 
 function decompress!(
@@ -719,74 +684,6 @@ function decompress!(
     return A
 end
 
-function decompress!(
-    A::SparseMatrixCSC{R},
-    Br::AbstractMatrix{R},
-    Bc::AbstractMatrix{R},
-    symmetric_to_row::Vector{Int},
-    symmetric_to_column::Vector{Int},
-    result::TreeSetColoringResult,
-) where {R<:Real}
-    (;
-        ag,
-        color,
-        reverse_bfs_orders,
-        diagonal_indices,
-        diagonal_nzind,
-        lower_triangle_offsets,
-        upper_triangle_offsets,
-        buffer,
-    ) = result
-    (; S) = ag
-    A_colptr = A.colptr
-    nzA = nonzeros(A)
-    m = size(Bc, 1)
-    n = size(Br, 2)
-
-    if eltype(buffer) == R
-        buffer_right_type = buffer
-    else
-        buffer_right_type = similar(buffer, R)
-    end
-
-    # Index of offsets in lower_triangle_offsets and upper_triangle_offsets
-    counter = 0
-
-    # Recover the off-diagonal coefficients of A
-    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
-            if i ≤ n
-                val = Br[symmetric_to_row[color[j]], i] - buffer_right_type[i]
-            else
-                val = Bc[i - n, symmetric_to_column[color[j]]] - buffer_right_type[i]
-            end
-            buffer_right_type[j] = buffer_right_type[j] + val
-
-            #! format: off
-            # A[i,j] is in the lower triangular part of A
-            if in_triangle(i, j, :L)
-                nzind = A_colptr[j + 1] - lower_triangle_offsets[counter]
-                nzA[nzind] = val
-            # A[i,j] is in the upper triangular part of A
-            else
-                nzind = A_colptr[i + 1] - lower_triangle_offsets[counter]
-                nzA[nzind] = val
-            end
-            #! format: on
-        end
-    end
-    return A
-end
-
 ## MatrixInverseColoringResult
 
 function decompress!(
@@ -861,14 +758,59 @@ function Base.getindex(B::JoinCompressed, i::Int, j::Int)
     end
 end
 
-function Base.getindex(B::JoinCompressed, k::Int)
-    dim = B.m + B.n
-    j, i = divrem(k - 1, dim)
-    return getindex(B, i + 1, j + 1)
+## StarSetBicoloringResult
+
+function decompress!(
+    A::AbstractMatrix,
+    Br::AbstractMatrix,
+    Bc::AbstractMatrix,
+    result::StarSetBicoloringResult,
+)
+    (; ag, symmetric_color, symmetric_to_row, symmetric_to_column, star_set) = result
+    (; star, hub, spokes) = star_set
+    (; S) = ag
+    fill!(A, zero(eltype(A)))
+
+    m, n = size(A)
+    for s in eachindex(hub, spokes)
+        j = abs(hub[s])
+        cj = symmetric_color[j]
+        for i in spokes[s]
+            if in_triangle(i, j, :L)
+                A[i - n, j] = Bc[i - n, symmetric_to_column[cj]]
+            else
+                A[j - n, i] = Br[symmetric_to_row[cj], i]
+            end
+        end
+    end
+    return A
+end
+
+function decompress!(
+    A::SparseMatrixCSC,
+    Br::AbstractMatrix,
+    Bc::AbstractMatrix,
+    result::StarSetBicoloringResult,
+)
+    (; ag, A_indices, compressed_indices, pos_Br) = result
+    (; S) = ag
+    nzA = nonzeros(A)
+    for k in 1:pos_Br
+        nzA[A_indices[k]] = Br[compressed_indices[k]]
+    end
+    for k in (pos_Br + 1):length(nzA)
+        nzA[A_indices[k]] = Bc[compressed_indices[k]]
+    end
+    return A
 end
 
+## TreeSetBicoloringResult
+
 function decompress!(
-    A::AbstractMatrix, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult
+    A::AbstractMatrix,
+    Br::AbstractMatrix,
+    Bc::AbstractMatrix,
+    result::TreeSetBicoloringResult,
 )
     (; symmetric_to_row, symmetric_to_column, symmetric_result) = result
     m, n = size(A)
@@ -879,15 +821,19 @@ function decompress!(
 end
 
 function decompress!(
-    A::SparseMatrixCSC, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult
+    A::SparseMatrixCSC,
+    Br::AbstractMatrix,
+    Bc::AbstractMatrix,
+    result::TreeSetBicoloringResult,
 )
     (;
         symmetric_to_row, symmetric_to_column, symmetric_result, large_colptr, large_rowval
     ) = result
     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, Bc, symmetric_to_row, symmetric_to_column, symmetric_result)
+    decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L)
     return A
 end

src/interface.jl

@@ -301,12 +301,15 @@ function _coloring(
 )
     A_and_Aᵀ = bidirectional_pattern(A; symmetric_pattern)
     ag = AdjacencyGraph(A_and_Aᵀ; has_diagonal=false)
-    color, star_set = star_coloring(ag, algo.order; postprocessing=algo.postprocessing)
+    symmetric_color, star_set = star_coloring(
+        ag, algo.order; postprocessing=algo.postprocessing
+    )
     if speed_setting isa WithResult
-        symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
-        return BicoloringResult(A, ag, symmetric_result)
+        return StarSetBicoloringResult(A, ag, symmetric_color, star_set)
     else
-        row_color, column_color, _ = remap_colors(color, maximum(color), size(A)...)
+        row_color, column_color, _ = remap_colors(
+            symmetric_color, maximum(symmetric_color), size(A)...
+        )
         return row_color, column_color
     end
 end
@@ -321,12 +324,16 @@ function _coloring(
 ) where {R}
     A_and_Aᵀ = bidirectional_pattern(A; symmetric_pattern)
     ag = AdjacencyGraph(A_and_Aᵀ; has_diagonal=false)
-    color, tree_set = acyclic_coloring(ag, algo.order; postprocessing=algo.postprocessing)
+    symmetric_color, tree_set = acyclic_coloring(
+        ag, algo.order; postprocessing=algo.postprocessing
+    )
     if speed_setting isa WithResult
-        symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R)
-        return BicoloringResult(A, ag, symmetric_result)
+        symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, symmetric_color, tree_set, R)
+        return TreeSetBicoloringResult(A, ag, symmetric_result)
     else
-        row_color, column_color, _ = remap_colors(color, maximum(color), size(A)...)
+        row_color, column_color, _ = remap_colors(
+            symmetric_color, maximum(symmetric_color), size(A)...
+        )
         return row_color, column_color
     end
 end