Decompression of acyclic coloring with SparseMatrixCSC

Author: amontoison

src/decompression.jl

@@ -514,9 +514,17 @@ function decompress!(
     result::TreeSetColoringResult,
     uplo::Symbol=:F,
 ) where {R<:Real}
-    @compat (; S, color, vertices_by_tree, reverse_bfs_orders,
-               diagonal_indices, lower_triangle_offsets,
-               upper_triangle_offsets, buffer) = result
+    @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)
 
@@ -536,31 +544,50 @@ function decompress!(
     elseif uplo == :U
         for i in diagonal_indices
             # A[i, i] is the last element in column i
-            nzind = A.colptr[i+1] - 1
+            nzind = A.colptr[i + 1] - 1
             nzA[nzind] = B[i, color[i]]
         end
     else  # uplo == :F
-        for i in diagonal_indices
-            # A[i, i] is the central element in column i
-            nzind = (A.colptr[i] + A.colptr[i+1] - 1) ÷ 2
+        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)
-                A[i, j] = val
+                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)
-                A[j, i] = val
+                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

src/result.jl

@@ -228,6 +228,7 @@ struct TreeSetColoringResult{M,R} <:
     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}
@@ -241,6 +242,7 @@ function TreeSetColoringResult(
 
     # Vector for the decompression of the diagonal coefficients
     diagonal_indices = Int[]
+    diagonal_nzind = Int[]
     ndiag = 0
 
     n = size(S, 1)
@@ -250,6 +252,7 @@ function TreeSetColoringResult(
             i = rv[k]
             if i == j
                 push!(diagonal_indices, i)
+                push!(diagonal_nzind, k)
                 ndiag += 1
             end
         end
@@ -366,27 +369,30 @@ function TreeSetColoringResult(
                     end
 
                     # Update lower_triangle_offsets and upper_triangle_offsets
-                    index_offsets += 1
                     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]
+                    index_ji = S.colptr[i] - 1 + searchsortedfirst(col_i, j)  # 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 the index (A.colptr[j+1] - offset_L) in A.nzval
-                        # offset_L = ...
-                        # lower_triangle_offsets[index_offsets] = S.colptr[j+1] - offset_L
+                        # 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 the index (A.colptr[i] + offset_U) in A.nzval
-                        # offset_U = ...
-                        # upper_triangle_offsets[index_offsets] = S.colptr[i] + offset_U
+                        # 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 the index (A.colptr[i+1] - offset_L) in A.nzval
-                        # offset_L = ...
-                        # lower_triangle_offsets[index_offsets] = S.colptr[i+1] - offset_L
+                        # 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 the index (A.colptr[j] + offset) in A.nzval
-                        # offset_U = ...
-                        # upper_triangle_offsets[index_offsets] = S.colptr[j] + offset_U
+                        # 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
@@ -398,8 +404,16 @@ function TreeSetColoringResult(
     buffer = Vector{R}(undef, nvertices)
 
     return TreeSetColoringResult(
-        S, color, group, vertices_by_tree, reverse_bfs_orders,
-        diagonal_indices, lower_triangle_offsets, upper_triangle_offsets, buffer
+        S,
+        color,
+        group,
+        vertices_by_tree,
+        reverse_bfs_orders,
+        diagonal_indices,
+        diagonal_nzind,
+        lower_triangle_offsets,
+        upper_triangle_offsets,
+        buffer,
     )
 end