Decompression of acyclic coloring with SparseMatrixCSC

Author: amontoison

src/decompression.jl

@@ -468,8 +468,6 @@ end
 
 ## TreeSetColoringResult
 
-# TODO: add method for A::SparseMatrixCSC
-
 function decompress!(
     A::AbstractMatrix{R},
     B::AbstractMatrix{R},
@@ -513,6 +511,92 @@ function decompress!(
     return A
 end
 
+function decompress!(
+    A::SparseMatrixCSC{R},
+    B::AbstractMatrix{R},
+    result::TreeSetColoringResult,
+    uplo::Symbol=:F,
+) where {R<:Real}
+    @compat (;
+        S,
+        color,
+        vertices_by_tree,
+        reverse_bfs_orders,
+        diagonal_indices,
+        diagonal_nzind,
+        lower_triangle_offsets,
+        upper_triangle_offsets,
+        buffer,
+    ) = result
+    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
+
+            # A[i,j] = val
+            if in_triangle(i, j, uplo)
+                if uplo == :L
+                    nzind = A.colptr[j + 1] - lower_triangle_offsets[counter]
+                else
+                    nzind = A.colptr[j] + upper_triangle_offsets[counter]
+                end
+                # Both values of `nzind` can be used if uplo = :F
+                nzA[nzind] = val
+            end
+
+            # A[j,i] = val
+            if in_triangle(j, i, uplo)
+                if uplo == :L
+                    nzind = A.colptr[i + 1] - lower_triangle_offsets[counter]
+                else
+                    nzind = A.colptr[i] + upper_triangle_offsets[counter]
+                end
+                # Both values of `nzind` can be used if uplo = :F
+                nzA[nzind] = val
+            end
+        end
+    end
+    return A
+end
+
 ## MatrixInverseColoringResult
 
 function decompress!(

src/result.jl

@@ -227,6 +227,10 @@ struct TreeSetColoringResult{M,R} <:
     group::Vector{Vector{Int}}
     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
 
@@ -236,6 +240,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)
@@ -297,6 +324,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]
 
@@ -317,7 +347,7 @@ function TreeSetColoringResult(
         end
 
         # continue until all leaves are treated
-        while queue_start <= queue_end
+        while queue_start ≤ queue_end
             leaf = queue[queue_start]
             queue_start += 1
 
@@ -337,6 +367,35 @@ 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_ij = S.colptr[j] - 1 + searchsortedfirst(col_j, i)  # S.nzval[index_ij] = S[i,j]
+                    @assert S.nzval[index_ij] == S[i,j]
+                    index_ji = S.colptr[i] - 1 + searchsortedfirst(col_i, j)  # S.nzval[index_ji] = S[j,i]
+                    @assert S.nzval[index_ji] == S[j,i]
+                    index_offsets += 1
+
+                    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
+                        lower_triangle_offsets[index_offsets] = S.colptr[j + 1] - index_ij
+
+                        # uplo = :U or uplo = :F
+                        # A[j,i] is stored at index_ji = (A.colptr[i] + offset_U) in A.nzval
+                        upper_triangle_offsets[index_offsets] = index_ji - S.colptr[i]
+                    else
+                        # uplo = :L or uplo = :F
+                        # A[j,i] is stored at index_ji = (A.colptr[i+1] - offset_L) in A.nzval
+                        lower_triangle_offsets[index_offsets] = S.colptr[i + 1] - index_ji
+
+                        # uplo = :U or uplo = :F
+                        # A[i,j] is stored at index_ij = (A.colptr[j] + offset_U) in A.nzval
+                        upper_triangle_offsets[index_offsets] = index_ij - S.colptr[j]
+                    end
                 end
             end
         end
@@ -347,7 +406,16 @@ function TreeSetColoringResult(
     buffer = Vector{R}(undef, nvertices)
 
     return TreeSetColoringResult(
-        S, color, group, vertices_by_tree, reverse_bfs_orders, buffer
+        S,
+        color,
+        group,
+        vertices_by_tree,
+        reverse_bfs_orders,
+        diagonal_indices,
+        diagonal_nzind,
+        lower_triangle_offsets,
+        upper_triangle_offsets,
+        buffer,
     )
 end