Add option to decompress into a triangle only (#83)

* Add option to decompress into a triangle only

* Add tests

Author: gdalle

src/decompression.jl

@@ -124,7 +124,7 @@ end
 """
     decompress!(
         A::AbstractMatrix, B::AbstractMatrix,
-        result::AbstractColoringResult,
+        result::AbstractColoringResult, [uplo=:UL]
     )
 
 Decompress `B` in-place into `A`, given a coloring `result` of the sparsity pattern of `A`.
@@ -133,6 +133,8 @@ The out-of-place alternative is [`decompress`](@ref).
 Compression means summing either the columns or the rows of `A` which share the same color.
 It is done by calling [`compress`](@ref).
 
+For `:symmetric` coloring results (and for those only), an optional positional argument `uplo in (:U, :L, :UL)` can be passed to specify which triangle of the matrix `A` should be updated: the upper one, the lower one, or both.
+
 !!! note
     In-place decompression is faster when `A isa SparseMatrixCSC`.
 
@@ -186,7 +188,7 @@ function decompress! end
 """
     decompress_single_color!(
         A::AbstractMatrix, b::AbstractVector, c::Integer,
-        result::AbstractColoringResult,
+        result::AbstractColoringResult, [uplo=:UL]
     )
 
 Decompress the vector `b` corresponding to color `c` in-place into `A`, given a coloring `result` of the sparsity pattern of `A`.
@@ -195,6 +197,8 @@ Decompress the vector `b` corresponding to color `c` in-place into `A`, given a
 - If `result` comes from a `:nonsymmetric` structure with `:row` partition, this will update the rows of `A` that share color `c` (whose sum makes up `b`).
 - If `result` comes from a `:symmetric` structure with `:column` partition, this will update the coefficients of `A` whose value is deduced from color `c`.
 
+For `:symmetric` coloring results (and for those only), an optional positional argument `uplo in (:U, :L, :UL)` can be passed to specify which triangle of the matrix `A` should be updated: the upper one, the lower one, or both.
+
 !!! warning
     This function will only update some coefficients of `A`, without resetting the rest to zero.
 
@@ -246,6 +250,16 @@ true
 """
 function decompress_single_color! end
 
+function in_triangle(i::Integer, j::Integer, uplo::Symbol)
+    if uplo == :UL
+        return true
+    elseif uplo == :U
+        return i <= j
+    else
+        return i >= j
+    end
+end
+
 ## ColumnColoringResult
 
 function decompress!(
@@ -293,6 +307,22 @@ function decompress!(
     return A
 end
 
+function decompress_single_color!(
+    A::SparseMatrixCSC{R}, b::AbstractVector{R}, c::Integer, result::ColumnColoringResult
+) where {R<:Real}
+    @compat (; S, group) = result
+    check_same_pattern(A, S)
+    rvS = rowvals(S)
+    nzA = nonzeros(A)
+    for j in group[c]
+        for k in nzrange(S, j)
+            i = rvS[k]
+            nzA[k] = b[i]
+        end
+    end
+    return A
+end
+
 ## RowColoringResult
 
 function decompress!(
@@ -343,7 +373,10 @@ end
 ## StarSetColoringResult
 
 function decompress!(
-    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::StarSetColoringResult
+    A::AbstractMatrix{R},
+    B::AbstractMatrix{R},
+    result::StarSetColoringResult,
+    uplo::Symbol=:UL,
 ) where {R<:Real}
     @compat (; S, color, star_set) = result
     @compat (; star, hub, spokes) = star_set
@@ -356,16 +389,25 @@ function decompress!(
     end
     for s in eachindex(hub, spokes)
         j = abs(hub[s])
+        cj = color[j]
         for i in spokes[s]
-            A[i, j] = B[i, color[j]]
-            A[j, i] = B[i, color[j]]
+            if in_triangle(i, j, uplo)
+                A[i, j] = B[i, cj]
+            end
+            if in_triangle(j, i, uplo)
+                A[j, i] = B[i, cj]
+            end
         end
     end
     return A
 end
 
 function decompress_single_color!(
-    A::AbstractMatrix{R}, b::AbstractVector{R}, c::Integer, result::StarSetColoringResult
+    A::AbstractMatrix{R},
+    b::AbstractVector{R},
+    c::Integer,
+    result::StarSetColoringResult,
+    uplo::Symbol=:UL,
 ) where {R<:Real}
     @compat (; S, color, group, star_set) = result
     @compat (; hub, spokes) = star_set
@@ -379,8 +421,12 @@ function decompress_single_color!(
         j = abs(hub[s])
         if color[j] == c
             for i in spokes[s]
-                A[i, j] = b[i]
-                A[j, i] = b[i]
+                if in_triangle(i, j, uplo)
+                    A[i, j] = b[i]
+                end
+                if in_triangle(j, i, uplo)
+                    A[j, i] = b[i]
+                end
             end
         end
     end
@@ -399,12 +445,33 @@ function decompress!(
     return A
 end
 
+function decompress!(
+    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::StarSetColoringResult, uplo::Symbol
+) where {R<:Real}
+    @compat (; S, compressed_indices) = result
+    check_same_pattern(A, S)
+    rvA = rowvals(A)
+    nzA = nonzeros(A)
+    for j in axes(S, 2)
+        for k in nzrange(S, j)
+            i = rvA[k]
+            if in_triangle(i, j, uplo)
+                nzA[k] = B[compressed_indices[k]]
+            end
+        end
+    end
+    return A
+end
+
 ## TreeSetColoringResult
 
 # TODO: add method for A::SparseMatrixCSC
 
 function decompress!(
-    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::TreeSetColoringResult
+    A::AbstractMatrix{R},
+    B::AbstractMatrix{R},
+    result::TreeSetColoringResult,
+    uplo::Symbol=:UL,
 ) where {R<:Real}
     @compat (; S, color, vertices_by_tree, reverse_bfs_orders, buffer) = result
     check_same_pattern(A, S)
@@ -432,8 +499,12 @@ function decompress!(
         for (i, j) in reverse_bfs_orders[k]
             val = B[i, color[j]] - buffer_right_type[i]
             buffer_right_type[j] = buffer_right_type[j] + val
-            A[i, j] = val
-            A[j, i] = val
+            if in_triangle(i, j, uplo)
+                A[i, j] = val
+            end
+            if in_triangle(j, i, uplo)
+                A[j, i] = val
+            end
         end
     end
     return A
@@ -442,7 +513,7 @@ end
 ## MatrixInverseColoringResult
 
 function decompress!(
-    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::LinearSystemColoringResult
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::LinearSystemColoringResult; uplo=:UL
 ) where {R<:Real}
     @compat (;
         S, color, strict_upper_nonzero_inds, T_factorization, strict_upper_nonzeros_A
@@ -458,8 +529,12 @@ function decompress!(
         end
     end
     for (l, (i, j)) in enumerate(strict_upper_nonzero_inds)
-        A[i, j] = strict_upper_nonzeros_A[l]
-        A[j, i] = strict_upper_nonzeros_A[l]
+        if in_triangle(i, j, uplo)
+            A[i, j] = strict_upper_nonzeros_A[l]
+        end
+        if in_triangle(j, i, uplo)
+            A[j, i] = strict_upper_nonzeros_A[l]
+        end
     end
     return A
 end

test/random.jl

@@ -1,6 +1,5 @@
 using ADTypes: column_coloring, row_coloring, symmetric_coloring
-using Compat
-using LinearAlgebra: I, Symmetric
+using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
 using SparseMatrixColorings:
@@ -27,7 +26,7 @@ symmetric_params = vcat(
 @testset "Column coloring & decompression" begin
     problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
     algo = GreedyColoringAlgorithm(; decompression=:direct)
-    for (m, n, p) in asymmetric_params
+    @testset "$((; m, n, p))" for (m, n, p) in asymmetric_params
         A0 = sprand(rng, m, n, p)
         color0 = column_coloring(A0, algo)
         @test structurally_orthogonal_columns(A0, color0)
@@ -39,7 +38,7 @@ end;
 @testset "Row coloring & decompression" begin
     problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
     algo = GreedyColoringAlgorithm(; decompression=:direct)
-    for (m, n, p) in asymmetric_params
+    @testset "$((; m, n, p))" for (m, n, p) in asymmetric_params
         A0 = sprand(rng, m, n, p)
         color0 = row_coloring(A0, algo)
         @test structurally_orthogonal_columns(transpose(A0), color0)
@@ -51,8 +50,8 @@ end;
 @testset "Symmetric coloring & direct decompression" begin
     problem = ColoringProblem(; structure=:symmetric, partition=:column)
     algo = GreedyColoringAlgorithm(; decompression=:direct)
-    for (n, p) in symmetric_params
-        A0 = Symmetric(sprand(rng, n, n, p))
+    @testset "$((; n, p))" for (n, p) in symmetric_params
+        A0 = sparse(Symmetric(sprand(rng, n, n, p)))
         color0 = symmetric_coloring(A0, algo)
         @test symmetrically_orthogonal_columns(A0, color0)
         @test directly_recoverable_columns(A0, color0)
@@ -63,8 +62,8 @@ end;
 @testset "Symmetric coloring & substitution decompression" begin
     problem = ColoringProblem(; structure=:symmetric, partition=:column)
     algo = GreedyColoringAlgorithm(; decompression=:substitution)
-    for (n, p) in symmetric_params
-        A0 = Symmetric(sprand(rng, n, n, p))
+    @testset "$((; n, p))" for (n, p) in symmetric_params
+        A0 = sparse(Symmetric(sprand(rng, n, n, p)))
         # TODO: find tests for recoverability
         test_coloring_decompression(A0, problem, algo)
     end

test/small.jl

@@ -1,5 +1,3 @@
-using Compat
-using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
 using SparseMatrixColorings:

test/utils.jl

@@ -1,3 +1,4 @@
+using LinearAlgebra
 using SparseMatrixColorings
 using SparseMatrixColorings: LinearSystemColoringResult, matrix_versions, respectful_similar
 using Test
@@ -10,38 +11,94 @@ function test_coloring_decompression(
     color0=nothing,
 ) where {structure,partition,decompression}
     color_vec = Vector{Int}[]
-    for A in matrix_versions(A0)
+    @testset "$(typeof(A))" for A in matrix_versions(A0)
         result = coloring(A, problem, algo; decompression_eltype=eltype(A))
         color = if partition == :column
             column_colors(result)
         elseif partition == :row
             row_colors(result)
         end
         push!(color_vec, color)
+
         B = compress(A, result)
-        !isnothing(color0) && @test color == color0
-        !isnothing(B0) && @test B == B0
-        @test decompress(B, result) ≈ A0
-        @test decompress(B, result) ≈ A0  # check result wasn't modified
-        @test decompress!(respectful_similar(A), B, result) ≈ A0
-        @test decompress!(respectful_similar(A), B, result) ≈ A0
-        if decompression == :direct
-            A2 = respectful_similar(A)
-            A2 .= zero(eltype(A2))
-            for c in unique(color)
-                if partition == :column
-                    decompress_single_color!(A2, B[:, c], c, result)
-                elseif partition == :row
-                    decompress_single_color!(A2, B[c, :], c, result)
+
+        @testset "Reference" begin
+            !isnothing(color0) && @test color == color0
+            !isnothing(B0) && @test B == B0
+        end
+
+        @testset "Full decompression" begin
+            @test decompress(B, result) ≈ A0
+            @test decompress(B, result) ≈ A0  # check result wasn't modified
+            @test decompress!(respectful_similar(A), B, result) ≈ A0
+            @test decompress!(respectful_similar(A), B, result) ≈ A0
+        end
+
+        @testset "Single-color decompression" begin
+            if decompression == :direct  # TODO: implement for :substitution too
+                A2 = respectful_similar(A)
+                A2 .= zero(eltype(A2))
+                for c in unique(color)
+                    if partition == :column
+                        decompress_single_color!(A2, B[:, c], c, result)
+                    elseif partition == :row
+                        decompress_single_color!(A2, B[c, :], c, result)
+                    end
                 end
+                @test A2 ≈ A0
             end
-            @test A2 ≈ A0
         end
-        if structure == :symmetric
-            linresult = LinearSystemColoringResult(sparse(A), color, eltype(A))
-            @test decompress(B, linresult) ≈ A0
-            @test decompress!(respectful_similar(A), B, linresult) ≈ A0
+
+        @testset "Triangle decompression" begin
+            if structure == :symmetric
+                A3upper = respectful_similar(A)
+                A3lower = respectful_similar(A)
+                A3both = respectful_similar(A)
+                A3upper .= zero(eltype(A))
+                A3lower .= zero(eltype(A))
+                A3both .= zero(eltype(A))
+
+                decompress!(A3upper, B, result, :U)
+                decompress!(A3lower, B, result, :L)
+                decompress!(A3both, B, result, :UL)
+
+                @test A3upper ≈ triu(A0)
+                @test A3lower ≈ tril(A0)
+                @test A3both ≈ A0
+            end
         end
+
+        @testset "Single-color triangle decompression" begin
+            if structure == :symmetric && decompression == :direct
+                A4upper = respectful_similar(A)
+                A4lower = respectful_similar(A)
+                A4both = respectful_similar(A)
+                A4upper .= zero(eltype(A))
+                A4lower .= zero(eltype(A))
+                A4both .= zero(eltype(A))
+
+                for c in unique(color)
+                    decompress_single_color!(A4upper, B[:, c], c, result, :U)
+                    decompress_single_color!(A4lower, B[:, c], c, result, :L)
+                    decompress_single_color!(A4both, B[:, c], c, result, :UL)
+                end
+
+                @test A4upper ≈ triu(A0)
+                @test A4lower ≈ tril(A0)
+                @test A4both ≈ A0
+            end
+        end
+
+        @testset "Linear system decompression" begin
+            if structure == :symmetric
+                linresult = LinearSystemColoringResult(sparse(A), color, eltype(A))
+                @test decompress(B, linresult) ≈ A0
+                @test decompress!(respectful_similar(A), B, linresult) ≈ A0
+            end
+        end
+    end
+
+    @testset "Coherence between all colorings" begin
+        @test all(color_vec .== Ref(color_vec[1]))
     end
-    @test all(color_vec .== Ref(color_vec[1]))
 end