Implement compress (#58)

* Implement `compress` as a counterpart to `decompress`

* Add doctests

Author: gdalle

docs/src/api.md

@@ -32,6 +32,7 @@ row_groups
 ## Decompression
 
 ```@docs
+compress
 decompress
 decompress!
 ```

src/SparseMatrixColorings.jl

@@ -48,12 +48,12 @@ include("check.jl")
 include("sparsematrixcsc.jl")
 include("examples.jl")
 
-@compat public NaturalOrder, RandomOrder, LargestFirst
-
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
 export coloring
 export column_colors, row_colors
 export column_groups, row_groups
-export decompress, decompress!
+export compress, decompress, decompress!
+
+@compat public NaturalOrder, RandomOrder, LargestFirst
 
 end

src/decompression.jl

@@ -1,7 +1,106 @@
+"""
+    compress(A, result::AbstractColoringResult)
+
+Compress `A` into a new matrix `B`, given a coloring `result` of the sparsity pattern of `A`.
+
+Compression means summing either the columns or the rows of `A` which share the same color.
+It is undone by calling [`decompress`](@ref).
+
+# Example
+
+```jldoctest
+julia> using SparseMatrixColorings, SparseArrays
+
+julia> A = sparse([
+           0 0 4 6 0 9
+           1 0 0 0 7 0
+           0 2 0 0 8 0
+           0 3 5 0 0 0
+       ]);
+
+julia> result = coloring(A, ColoringProblem(), GreedyColoringAlgorithm());
+
+julia> column_groups(result)
+3-element Vector{Vector{Int64}}:
+ [1, 2, 4]
+ [3, 5]
+ [6]
+
+julia> B = compress(A, result)
+4×3 Matrix{Int64}:
+ 6  4  9
+ 1  7  0
+ 2  8  0
+ 3  5  0
+```
+
+# See also
+
+- [`AbstractColoringResult`](@ref)
+"""
+function compress end
+
+function compress(A, result::AbstractColoringResult{structure,:column}) where {structure}
+    group = column_groups(result)
+    B = stack(group; dims=2) do g
+        dropdims(sum(A[:, g]; dims=2); dims=2)
+    end
+    return B
+end
+
+function compress(A, result::AbstractColoringResult{structure,:row}) where {structure}
+    group = row_groups(result)
+    B = stack(group; dims=1) do g
+        dropdims(sum(A[g, :]; dims=1); dims=1)
+    end
+    return B
+end
+
 """
     decompress(B::AbstractMatrix, result::AbstractColoringResult)
 
-Decompress `B` out-of-place into a new matrix `A`, given a coloring `result` of the sparsity pattern of `A`.
+Decompress `B` into a new matrix `A`, given a coloring `result` of the sparsity pattern of `A`.
+
+Compression means summing either the columns or the rows of `A` which share the same color.
+It is done by calling [`compress`](@ref).
+
+# Example
+
+```jldoctest
+julia> using SparseMatrixColorings, SparseArrays
+
+julia> A = sparse([
+           0 0 4 6 0 9
+           1 0 0 0 7 0
+           0 2 0 0 8 0
+           0 3 5 0 0 0
+       ]);
+
+julia> result = coloring(A, ColoringProblem(), GreedyColoringAlgorithm());
+
+julia> column_groups(result)
+3-element Vector{Vector{Int64}}:
+ [1, 2, 4]
+ [3, 5]
+ [6]
+
+julia> B = compress(A, result)
+4×3 Matrix{Int64}:
+ 6  4  9
+ 1  7  0
+ 2  8  0
+ 3  5  0
+
+julia> decompress(B, result)
+4×6 SparseMatrixCSC{Int64, Int64} with 9 stored entries:
+ ⋅  ⋅  4  6  ⋅  9
+ 1  ⋅  ⋅  ⋅  7  ⋅
+ ⋅  2  ⋅  ⋅  8  ⋅
+ ⋅  3  5  ⋅  ⋅  ⋅
+
+julia> decompress(B, result) == A
+true
+```
 
 # See also
 
@@ -19,7 +118,50 @@ end
         result::AbstractColoringResult,
     )
 
-Decompress `B` in-place into an existing matrix `A`, given a coloring `result` of the sparsity pattern of `A`.
+Decompress `B` in-place into `A`, given a coloring `result` of the sparsity pattern of `A`.
+
+Compression means summing either the columns or the rows of `A` which share the same color.
+It is done by calling [`compress`](@ref).
+
+# Example
+
+```jldoctest
+julia> using SparseMatrixColorings, SparseArrays
+
+julia> A = sparse([
+           0 0 4 6 0 9
+           1 0 0 0 7 0
+           0 2 0 0 8 0
+           0 3 5 0 0 0
+       ]);
+
+julia> result = coloring(A, ColoringProblem(), GreedyColoringAlgorithm());
+
+julia> column_groups(result)
+3-element Vector{Vector{Int64}}:
+ [1, 2, 4]
+ [3, 5]
+ [6]
+
+julia> B = compress(A, result)
+4×3 Matrix{Int64}:
+ 6  4  9
+ 1  7  0
+ 2  8  0
+ 3  5  0
+
+julia> A2 = similar(A);
+
+julia> decompress!(A2, B, result)
+4×6 SparseMatrixCSC{Int64, Int64} with 9 stored entries:
+ ⋅  ⋅  4  6  ⋅  9
+ 1  ⋅  ⋅  ⋅  7  ⋅
+ ⋅  2  ⋅  ⋅  8  ⋅
+ ⋅  3  5  ⋅  ⋅  ⋅
+
+julia> A2 == A
+true
+```
 
 # See also
 

test/random.jl

@@ -38,10 +38,7 @@ symmetric_params = vcat(
         @testset "A::$(typeof(A))" for A in matrix_versions(A0)
             result = coloring(A, problem, algo)
             color = column_colors(result)
-            group = column_groups(result)
-            B = stack(group; dims=2) do g
-                dropdims(sum(A[:, g]; dims=2); dims=2)
-            end
+            B = compress(A, result)
             @test color == column_coloring(A, algo)
             @test structurally_orthogonal_columns(A, color)
             @test directly_recoverable_columns(A, color)
@@ -60,10 +57,7 @@ end;
         @testset "A::$(typeof(A))" for A in matrix_versions(A0)
             result = coloring(A, problem, algo)
             color = row_colors(result)
-            group = row_groups(result)
-            B = stack(group; dims=1) do g
-                dropdims(sum(A[g, :]; dims=1); dims=1)
-            end
+            B = compress(A, result)
             @test color == row_coloring(A, algo)
             @test structurally_orthogonal_columns(transpose(A), color)
             @test directly_recoverable_columns(transpose(A), color)
@@ -88,10 +82,7 @@ end;
             @testset "A::$(typeof(A))" for A in matrix_versions(A0)
                 result = coloring(A, problem, algo)
                 color = column_colors(result)
-                group = column_groups(result)
-                B = stack(group; dims=2) do g
-                    dropdims(sum(A[:, g]; dims=2); dims=2)
-                end
+                B = compress(A, result)
                 if key == :direct
                     @test color == symmetric_coloring(A, algo)
                     @test symmetrically_orthogonal_columns(A, color)

test/small.jl

@@ -35,6 +35,7 @@ algo = GreedyColoringAlgorithm()
     result0 = DefaultColoringResult{:nonsymmetric,:column,:direct}(A0, color0)
     @test structurally_orthogonal_columns(A0, color0)
     @test directly_recoverable_columns(A0, color0)
+    @test compress(A0, result0) == B0
     @test decompress(B0, result0) == A0
     for A in matrix_versions(A0)
         @test decompress!(respectful_similar(A), B0, result0) == A
@@ -55,6 +56,7 @@ end;
     result0 = DefaultColoringResult{:nonsymmetric,:row,:direct}(A0, color0)
     @test structurally_orthogonal_columns(transpose(A0), color0)
     @test directly_recoverable_columns(transpose(A0), color0)
+    @test compress(A0, result0) == B0
     @test decompress(B0, result0) == A0
     for A in matrix_versions(A0)
         @test decompress!(respectful_similar(A), B0, result0) == A
@@ -68,6 +70,7 @@ end;
         result0 = DefaultColoringResult{:symmetric,:column,:direct}(A0, color0)
         @test symmetrically_orthogonal_columns(A0, color0)
         @test directly_recoverable_columns(A0, color0)
+        @test compress(A0, result0) == B0
         @test decompress(B0, result0) == A0
         for A in matrix_versions(A0)
             @test decompress!(respectful_similar(A), B0, result0) == A
@@ -84,13 +87,11 @@ end;
                 structure=:symmetric, partition=:column, decompression=:substitution
             ),
             GreedyColoringAlgorithm(),
-        )  # returns a TreeSetColoringResult
-        group = column_groups(result)
-        B = stack(group; dims=2) do g
-            dropdims(sum(A0[:, g]; dims=2); dims=2)
-        end
+        )
+        B = compress(A0, result)
         @test column_colors(result) != color0
         @test B != B0
+        @test compress(A0, result0) == B0
         @test decompress(B, result) ≈ A0
         @test decompress(B0, result0) ≈ A0
         for A in matrix_versions(A0)
@@ -105,6 +106,7 @@ end;
         result0 = DefaultColoringResult{:symmetric,:column,:direct}(A0, color0)
         @test symmetrically_orthogonal_columns(A0, color0)
         @test directly_recoverable_columns(A0, color0)
+        @test compress(A0, result0) == B0
         @test decompress(B0, result0) == A0
         for A in matrix_versions(A0)
             @test decompress!(respectful_similar(A), B0, result0) == A
@@ -121,8 +123,10 @@ end;
                 structure=:symmetric, partition=:column, decompression=:substitution
             ),
             GreedyColoringAlgorithm(),
-        )  # returns a TreeSetColoringResult
+        )
         @test column_colors(result) == color0
+        @test compress(A0, result0) == B0
+        @test compress(A0, result) == B0
         @test decompress(B0, result0) ≈ A0
         @test decompress(B0, result) ≈ A0
         for A in matrix_versions(A0)