Add an option decompression_uplo for symmetric results

Author: amontoison

src/decompression.jl

@@ -453,12 +453,18 @@ function decompress!(
     uplo == :F && check_same_pattern(A, S)
     fill!(A, zero(eltype(A)))
 
+    l = 0
     rvS = rowvals(S)
     for j in axes(S, 2)
         for k in nzrange(S, j)
             i = rvS[k]
             if in_triangle(i, j, uplo)
-                A[i, j] = B[compressed_indices[k]]
+                if result.decompression_uplo == :F
+                    A[i, j] = B[compressed_indices[k]]
+                else
+                    l += 1
+                    A[i, j] = B[compressed_indices[l]]
+                end
             end
         end
     end
@@ -472,6 +478,7 @@ function decompress_single_color!(
     result::StarSetColoringResult,
     uplo::Symbol=:F,
 )
+    @assert result.decompression_uplo == :F
     (; ag, compressed_indices, group) = result
     (; S) = ag
     uplo == :F && check_same_pattern(A, S)
@@ -508,12 +515,13 @@ function decompress!(
     (; ag, compressed_indices) = result
     (; S) = ag
     nzA = nonzeros(A)
-    if uplo == :F
-        check_same_pattern(A, S)
+    if result.decompression_uplo == uplo
+        uplo == :F && check_same_pattern(A, S)
         for k in eachindex(nzA, compressed_indices)
             nzA[k] = B[compressed_indices[k]]
         end
     else
+        @assert result.decompression_uplo == :F
         rvS = rowvals(S)
         l = 0  # assume A has the same pattern as the triangle
         for j in axes(S, 2)

src/interface.jl

@@ -345,7 +345,9 @@ function _coloring(
         t -> maximum(t[3]) + maximum(t[4]), outputs_by_order
     )  # can't use ncolors without computing the full result
     if speed_setting isa WithResult
-        symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
+        symmetric_result = StarSetColoringResult(
+            A_and_Aᵀ, ag, color, star_set; decompression_uplo=:L
+        )
         return BicoloringResult(
             A,
             ag,
@@ -390,7 +392,9 @@ function _coloring(
         t -> maximum(t[3]) + maximum(t[4]), outputs_by_order
     )  # can't use ncolors without computing the full result
     if speed_setting isa WithResult
-        symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R)
+        symmetric_result = TreeSetColoringResult(
+            A_and_Aᵀ, ag, color, tree_set, R; decompression_uplo=:L
+        )
         return BicoloringResult(
             A,
             ag,

src/result.jl

@@ -309,55 +309,74 @@ struct StarSetColoringResult{
     color::CT
     group::GT
     compressed_indices::VT
+    decompression_uplo::Symbol
     additional_info::A
 end
 
 function StarSetColoringResult(
     A::AbstractMatrix,
     ag::AdjacencyGraph{T},
     color::Vector{<:Integer},
-    star_set::StarSet{<:Integer},
+    star_set::StarSet{<:Integer};
+    decompression_uplo::Symbol=:F,
 ) where {T<:Integer}
     group = group_by_color(T, color)
-    compressed_indices = star_csc_indices(ag, color, star_set)
-    return StarSetColoringResult(A, ag, color, group, compressed_indices, nothing)
+    compressed_indices = star_csc_indices(ag, color, star_set, decompression_uplo)
+    return StarSetColoringResult(
+        A, ag, color, group, compressed_indices, decompression_uplo, nothing
+    )
 end
 
 function star_csc_indices(
-    ag::AdjacencyGraph{T}, color::Vector{<:Integer}, star_set
+    ag::AdjacencyGraph{T},
+    color::Vector{<:Integer},
+    star_set::StarSet{<:Integer},
+    decompression_uplo::Symbol,
 ) where {T}
     (; star, hub) = star_set
     S = pattern(ag)
     edge_to_index = edge_indices(ag)
     n = S.n
     rvS = rowvals(S)
-    compressed_indices = zeros(T, nnz(S))  # needs to be independent from the storage in the graph, in case the graph gets reused
+    l = 0
+    if augmented_graph(ag) && (decompression_uplo != :F)
+        compressed_indices = zeros(T, nnz(S) ÷ 2)
+    else
+        compressed_indices = zeros(T, nnz(S))  # needs to be independent from the storage in the graph, in case the graph gets reused
+    end
     for j in axes(S, 2)
         for k in nzrange(S, j)
             i = rvS[k]
             if i == j
                 # diagonal coefficients
+                l += 1
                 c = color[i]
-                compressed_indices[k] = (c - 1) * n + i
+                compressed_indices[l] = (c - 1) * n + i
             else
-                # off-diagonal coefficients
-                index_ij = edge_to_index[k]
-                s = star[index_ij]
-                h = abs(hub[s])
-
-                # Assign the non-hub vertex (spoke) to the correct position in spokes
-                if i == h
-                    # i is the hub and j is the spoke
-                    c = color[i]
-                    compressed_indices[k] = (c - 1) * n + j
-                else  # j == h
-                    # j is the hub and i is the spoke
-                    c = color[j]
-                    compressed_indices[k] = (c - 1) * n + i
+                if in_triangle(i, j, decompression_uplo)
+                    # off-diagonal coefficients
+                    l += 1
+                    index_ij = edge_to_index[k]
+                    s = star[index_ij]
+                    h = abs(hub[s])
+
+                    # Assign the non-hub vertex (spoke) to the correct position in spokes
+                    if i == h
+                        # i is the hub and j is the spoke
+                        c = color[i]
+                        compressed_indices[l] = (c - 1) * n + j
+                    else  # j == h
+                        # j is the hub and i is the spoke
+                        c = color[j]
+                        compressed_indices[l] = (c - 1) * n + i
+                    end
                 end
             end
         end
     end
+    if !augmented_graph(ag) && (decompression_uplo != :F)
+        resize!(compressed_indices, l)
+    end
     return compressed_indices
 end
 
@@ -391,43 +410,59 @@ struct TreeSetColoringResult{
     lower_triangle_offsets::Vector{T}
     upper_triangle_offsets::Vector{T}
     buffer::Vector{R}
+    decompression_uplo::Symbol
 end
 
 function TreeSetColoringResult(
     A::AbstractMatrix,
     ag::AdjacencyGraph{T},
     color::Vector{<:Integer},
     tree_set::TreeSet{<:Integer},
-    decompression_eltype::Type{R},
+    decompression_eltype::Type{R};
+    decompression_uplo::Symbol=:F,
 ) where {T<:Integer,R}
     (; reverse_bfs_orders, tree_edge_indices, nt) = tree_set
     (; S) = ag
     nvertices = length(color)
     group = group_by_color(T, color)
     rv = rowvals(S)
 
-    # Vector for the decompression of the diagonal coefficients
-    diagonal_indices = T[]
-    diagonal_nzind = T[]
     ndiag = 0
-
     if !augmented_graph(ag)
         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
     end
 
+    # Vector for the decompression of the diagonal coefficients
+    diagonal_indices = Vector{T}(undef, ndiag)
+    diagonal_nzind = (decompression_uplo == :F) ? Vector{T}(undef, ndiag) : T[]
+
+    if !augmented_graph(ag)
+        l = 0
+        for j in axes(S, 2)
+            for k in nzrange(S, j)
+                i = rv[k]
+                if i == j
+                    l += 1
+                    diagonal_indices[l] = i
+                    if decompression_uplo == :F
+                        diagonal_nzind[l] = k
+                    end
+                end
+            end
+        end
+    end
+
     # Vectors for the decompression of the off-diagonal coefficients
     nedges = (nnz(S) - ndiag) ÷ 2
-    lower_triangle_offsets = Vector{T}(undef, nedges)
-    upper_triangle_offsets = Vector{T}(undef, nedges)
+    lower_triangle_offsets = decompression_uplo == :U ? T[] : Vector{T}(undef, nedges)
+    upper_triangle_offsets = decompression_uplo == :L ? T[] : Vector{T}(undef, nedges)
 
     # Index in lower_triangle_offsets and upper_triangle_offsets
     index_offsets = 0
@@ -451,21 +486,29 @@ function TreeSetColoringResult(
             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
+                if decompression_uplo != :U
+                    lower_triangle_offsets[index_offsets] = length(col_j) - searchsortedfirst(col_j, i) + 1
+                end
 
                 # 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
+                if decompression_uplo != :L
+                    upper_triangle_offsets[index_offsets] = searchsortedfirst(col_i, j)::Int - 1
+                end
 
             # 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
+                if decompression_uplo != :L
+                    upper_triangle_offsets[index_offsets] = searchsortedfirst(col_j, i)::Int - 1
+                end
 
                 # 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
+                if decompression_uplo != :U
+                    lower_triangle_offsets[index_offsets] = length(col_i) - searchsortedfirst(col_i, j) + 1
+                end
             end
             #! format: on
         end
@@ -488,6 +531,7 @@ function TreeSetColoringResult(
         lower_triangle_offsets,
         upper_triangle_offsets,
         buffer,
+        decompression_uplo,
     )
 end