Decompression of acyclic coloring with SparseMatrixCSC

src/decompression.jl

@@ -460,8 +460,6 @@ end
 
 ## TreeSetColoringResult
 
-# TODO: add method for A::SparseMatrixCSC
-
 function decompress!(
     A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
 )
@@ -504,6 +502,106 @@ function decompress!(
     return A
 end
 
+function decompress!(
+    A::SparseMatrixCSC{R},
+    B::AbstractMatrix{R},
+    result::TreeSetColoringResult,
+    uplo::Symbol=:F,
+) where {R<:Real}
+    (;
+        color,
+        vertices_by_tree,
+        reverse_bfs_orders,
+        diagonal_indices,
+        diagonal_nzind,
+        lower_triangle_offsets,
+        upper_triangle_offsets,
+        buffer,
+    ) = result
+    S = result.ag.S
+    A_colptr = A.colptr
+    nzA = nonzeros(A)
+    uplo == :F && check_same_pattern(A, S)
+
+    if eltype(buffer) == R
+        buffer_right_type = buffer
+    else
+        buffer_right_type = similar(buffer, R)
+    end
+
+    # Recover the diagonal coefficients of A
+    if uplo == :L
+        for i in diagonal_indices
+            # A[i, i] is the first element in column i
+            nzind = A_colptr[i]
+            nzA[nzind] = B[i, color[i]]
+        end
+    elseif uplo == :U
+        for i in diagonal_indices
+            # A[i, i] is the last element in column i
+            nzind = A_colptr[i + 1] - 1
+            nzA[nzind] = B[i, color[i]]
+        end
+    else  # uplo == :F
+        for (k, i) in enumerate(diagonal_indices)
+            nzind = diagonal_nzind[k]
+            nzA[nzind] = B[i, color[i]]
+        end
+    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(vertices_by_tree, reverse_bfs_orders)
+        for vertex in vertices_by_tree[k]
+            buffer_right_type[vertex] = zero(R)
+        end
+
+        for (i, j) in reverse_bfs_orders[k]
+            counter += 1
+            val = B[i, color[j]] - buffer_right_type[i]
+            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)
+                # uplo = :L or uplo = :F
+                # A[i,j] is stored at index_ij = (A.colptr[j+1] - offset_L) in A.nzval
+                if uplo != :U
+                    nzind = A_colptr[j + 1] - lower_triangle_offsets[counter]
+                    nzA[nzind] = val
+                end
+
+                # uplo = :U or uplo = :F
+                # A[j,i] is stored at index_ji = (A.colptr[i] + offset_U) in A.nzval
+                if uplo != :L
+                    nzind = A_colptr[i] + upper_triangle_offsets[counter]
+                    nzA[nzind] = val
+                end
+
+            # A[i,j] is in the upper triangular part of A
+            else
+                # uplo = :U or uplo = :F
+                # A[i,j] is stored at index_ij = (A.colptr[j] + offset_U) in A.nzval
+                if uplo != :L
+                    nzind = A_colptr[j] + upper_triangle_offsets[counter]
+                    nzA[nzind] = val
+                end
+
+                # uplo = :L or uplo = :F
+                # A[j,i] is stored at index_ji = (A.colptr[i+1] - offset_L) in A.nzval
+                if uplo != :U
+                    nzind = A_colptr[i + 1] - lower_triangle_offsets[counter]
+                    nzA[nzind] = val
+                end
+            end
+            #! format: on
+        end
+    end
+    return A
+end
+
 ## MatrixInverseColoringResult
 
 function decompress!(

src/result.jl

@@ -258,6 +258,10 @@ struct TreeSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V,R} <:
     group::V
     vertices_by_tree::Vector{Vector{Int}}
     reverse_bfs_orders::Vector{Vector{Tuple{Int,Int}}}
+    diagonal_indices::Vector{Int}
+    diagonal_nzind::Vector{Int}
+    lower_triangle_offsets::Vector{Int}
+    upper_triangle_offsets::Vector{Int}
     buffer::Vector{R}
 end
 
@@ -272,6 +276,29 @@ function TreeSetColoringResult(
     nvertices = length(color)
     group = group_by_color(color)
 
+    # Vector for the decompression of the diagonal coefficients
+    diagonal_indices = Int[]
+    diagonal_nzind = Int[]
+    ndiag = 0
+
+    n = size(S, 1)
+    rv = rowvals(S)
+    for j in axes(S, 2)
+        for k in nzrange(S, j)
+            i = rv[k]
+            if i == j
+                push!(diagonal_indices, i)
+                push!(diagonal_nzind, k)
+                ndiag += 1
+            end
+        end
+    end
+
+    # Vectors for the decompression of the off-diagonal coefficients
+    nedges = (nnz(S) - ndiag) ÷ 2
+    lower_triangle_offsets = Vector{Int}(undef, nedges)
+    upper_triangle_offsets = Vector{Int}(undef, nedges)
+
     # forest is a structure DisjointSets from DataStructures.jl
     # - forest.intmap: a dictionary that maps an edge (i, j) to an integer k
     # - forest.revmap: a dictionary that does the reverse of intmap, mapping an integer k to an edge (i, j)
@@ -333,6 +360,9 @@ function TreeSetColoringResult(
     # Create a queue with a fixed size nvmax
     queue = Vector{Int}(undef, nvmax)
 
+    # Index in lower_triangle_offsets and upper_triangle_offsets
+    index_offsets = 0
+
     for k in 1:ntrees
         tree = trees[k]
 
@@ -373,6 +403,36 @@ function TreeSetColoringResult(
                         queue_end += 1
                         queue[queue_end] = neighbor
                     end
+
+                    # Update lower_triangle_offsets and upper_triangle_offsets
+                    i = leaf
+                    j = neighbor
+                    col_i = view(rv, nzrange(S, i))
+                    col_j = view(rv, nzrange(S, j))
+                    index_offsets += 1
+
+                    #! format: off
+                    # S[i,j] is in the lower triangular part of S
+                    if in_triangle(i, j, :L)
+                        # uplo = :L or uplo = :F
+                        # S[i,j] is stored at index_ij = (S.colptr[j+1] - offset_L) in S.nzval
+                        lower_triangle_offsets[index_offsets] = length(col_j) - searchsortedfirst(col_j, i) + 1
+
+                        # uplo = :U or uplo = :F
+                        # S[j,i] is stored at index_ji = (S.colptr[i] + offset_U) in S.nzval
+                        upper_triangle_offsets[index_offsets] = searchsortedfirst(col_i, j)::Int - 1
+
+                    # S[i,j] is in the upper triangular part of S
+                    else
+                        # uplo = :U or uplo = :F
+                        # S[i,j] is stored at index_ij = (S.colptr[j] + offset_U) in S.nzval
+                        upper_triangle_offsets[index_offsets] = searchsortedfirst(col_j, i)::Int - 1
+
+                        # uplo = :L or uplo = :F
+                        # S[j,i] is stored at index_ji = (S.colptr[i+1] - offset_L) in S.nzval
+                        lower_triangle_offsets[index_offsets] = length(col_i) - searchsortedfirst(col_i, j) + 1
+                    end
+                    #! format: on
                 end
             end
         end
@@ -383,7 +443,17 @@ function TreeSetColoringResult(
     buffer = Vector{R}(undef, nvertices)
 
     return TreeSetColoringResult(
-        A, ag, color, group, vertices_by_tree, reverse_bfs_orders, buffer
+        A,
+        ag,
+        color,
+        group,
+        vertices_by_tree,
+        reverse_bfs_orders,
+        diagonal_indices,
+        diagonal_nzind,
+        lower_triangle_offsets,
+        upper_triangle_offsets,
+        buffer,
     )
 end