Fix star decompression into SparseMatrixCSC triangle (#90)

* Fix star decompression into `SparseMatrixCSC` triangle

* Test with `triu`

* Fix

* Check that if I remove l tests fail

* Document

* Add assertion

* Precise error

* Test throws

* Revert to assert

Author: gdalle

src/decompression.jl

@@ -130,13 +130,14 @@ end
 Decompress `B` in-place into `A`, given a coloring `result` of the sparsity pattern of `A`.
 The out-of-place alternative is [`decompress`](@ref).
 
+!!! note
+    In-place decompression is faster when `A isa SparseMatrixCSC`.
+
 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, :F)` can be passed to specify which part of the matrix `A` should be updated: the Upper triangle, the Lower triangle, or the Full matrix.
-
-!!! note
-    In-place decompression is faster when `A isa SparseMatrixCSC`.
+When `A isa SparseMatrixCSC`, using the `uplo` argument requires a target matrix which only stores the relevant triangle(s).
 
 # Example
 
@@ -197,11 +198,12 @@ 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, :F)` can be passed to specify which part of the matrix `A` should be updated: the Upper triangle, the Lower triangle, or the Full matrix.
-
 !!! warning
     This function will only update some coefficients of `A`, without resetting the rest to zero.
 
+For `:symmetric` coloring results (and for those only), an optional positional argument `uplo in (:U, :L, :F)` can be passed to specify which part of the matrix `A` should be updated: the Upper triangle, the Lower triangle, or the Full matrix.
+When `A isa SparseMatrixCSC`, using the `uplo` argument requires a target matrix which only stores the relevant triangle(s).
+
 # Example
 
 ```jldoctest
@@ -448,16 +450,19 @@ function decompress!(
 ) where {R<:Real}
     @compat (; S, compressed_indices) = result
     uplo == :F && check_same_pattern(A, S)
-    rvA = rowvals(A)
+    rvS = rowvals(S)
     nzA = nonzeros(A)
+    l = 0  # assume A has the same pattern as the triangle
     for j in axes(S, 2)
         for k in nzrange(S, j)
-            i = rvA[k]
+            i = rvS[k]
             if in_triangle(i, j, uplo)
-                nzA[k] = B[compressed_indices[k]]
+                l += 1
+                nzA[l] = B[compressed_indices[k]]
             end
         end
     end
+    @assert l == length(nonzeros(A))
     return A
 end
 

test/allocations.jl

@@ -54,7 +54,7 @@ function test_noallocs_decompression(
     end
     @testset "Triangle decompression" begin
         if structure == :symmetric
-            bench1_triangle = @be similar(A) decompress!(_, B, result, :U) evals = 1
+            bench1_triangle = @be similar(triu(A)) decompress!(_, B, result, :U) evals = 1
             bench2_triangle = @be similar(Matrix(A)) decompress!(_, B, result, :U) evals = 1
             @test minimum(bench1_triangle).allocs == 0
             @test minimum(bench2_triangle).allocs == 0
@@ -63,7 +63,7 @@ function test_noallocs_decompression(
     @testset "Single-color triangle decompression" begin
         if structure == :symmetric && decompression == :direct
             b = B[:, 1]
-            bench1_singlecolor_triangle = @be similar(A) decompress_single_color!(
+            bench1_singlecolor_triangle = @be similar(triu(A)) decompress_single_color!(
                 _, b, 1, result, :U
             ) evals = 1
             bench2_singlecolor_triangle = @be similar(Matrix(A)) decompress_single_color!(

test/type_stability.jl

@@ -79,13 +79,13 @@ end
             end
             @testset "Triangle decompression" begin
                 if structure == :symmetric
-                    @test_opt decompress!(respectful_similar(A), B, result, :U)
+                    @test_opt decompress!(respectful_similar(triu(A)), B, result, :U)
                 end
             end
             @testset "Single-color triangle decompression" begin
                 if structure == :symmetric && decompression == :direct
                     @test_opt decompress_single_color!(
-                        respectful_similar(A), B[:, 1], 1, result, :U
+                        respectful_similar(triu(A)), B[:, 1], 1, result, :U
                     )
                 end
             end

test/utils.jl

@@ -51,8 +51,8 @@ function test_coloring_decompression(
 
         @testset "Triangle decompression" begin
             if structure == :symmetric
-                A3upper = respectful_similar(A)
-                A3lower = respectful_similar(A)
+                A3upper = respectful_similar(triu(A))
+                A3lower = respectful_similar(tril(A))
                 A3both = respectful_similar(A)
                 A3upper .= zero(eltype(A))
                 A3lower .= zero(eltype(A))
@@ -70,8 +70,8 @@ function test_coloring_decompression(
 
         @testset "Single-color triangle decompression" begin
             if structure == :symmetric && decompression == :direct
-                A4upper = respectful_similar(A)
-                A4lower = respectful_similar(A)
+                A4upper = respectful_similar(triu(A))
+                A4lower = respectful_similar(tril(A))
                 A4both = respectful_similar(A)
                 A4upper .= zero(eltype(A))
                 A4lower .= zero(eltype(A))