Implement StarSetColoringResult and TreeSetColoringResult (#57)

* StarSet-based decompression

* Add result types to doc

* Add tests of standard and star/tree-set decompression

* Remove useless getters

* Docstrings

Author: gdalle

docs/src/dev.md

@@ -35,6 +35,8 @@ SparseMatrixColorings.TreeSet
 
 ```@docs
 SparseMatrixColorings.DefaultColoringResult
+SparseMatrixColorings.StarSetColoringResult
+SparseMatrixColorings.TreeSetColoringResult
 SparseMatrixColorings.DirectSparseColoringResult
 ```
 

src/decompression.jl

@@ -92,6 +92,20 @@ function decompress_aux!(
     return A
 end
 
+function decompress_aux!(
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::StarSetColoringResult{:column}
+) where {R<:Real}
+    A .= zero(R)
+    S = get_matrix(result)
+    color = column_colors(result)
+    for ij in findall(!iszero, S)
+        i, j = Tuple(ij)
+        k, l = symmetric_coefficient(i, j, color, result.star_set)
+        A[i, j] = B[k, l]
+    end
+    return A
+end
+
 function decompress_aux!(
     A::AbstractMatrix{R},
     B::AbstractMatrix{R},

src/interface.jl

@@ -150,8 +150,7 @@ function coloring(
 )
     ag = adjacency_graph(S)
     color, star_set = star_coloring(ag, algo.order)
-    # TODO: handle star_set
-    return DefaultColoringResult{:symmetric,:column,:direct}(S, color)
+    return StarSetColoringResult{:column}(S, color, star_set)
 end
 
 function coloring(
@@ -161,8 +160,7 @@ function coloring(
 )
     ag = adjacency_graph(S)
     color, tree_set = acyclic_coloring(ag, algo.order)
-    # TODO: handle tree_set
-    return DefaultColoringResult{:symmetric,:column,:substitution}(S, color)
+    return TreeSetColoringResult{:column}(S, color, tree_set)
 end
 
 ## ADTypes interface

src/result.jl

@@ -70,6 +70,14 @@ function group_by_color(color::AbstractVector{<:Integer})
     return group
 end
 
+get_matrix(result::AbstractColoringResult) = result.matrix
+
+column_colors(result::AbstractColoringResult{s,:column}) where {s} = result.color
+column_groups(result::AbstractColoringResult{s,:column}) where {s} = result.group
+
+row_colors(result::AbstractColoringResult{s,:row}) where {s} = result.color
+row_groups(result::AbstractColoringResult{s,:row}) where {s} = result.group
+
 ## Concrete subtypes
 
 """
@@ -103,10 +111,64 @@ function DefaultColoringResult{structure,partition,decompression}(
     )
 end
 
-get_matrix(result::DefaultColoringResult) = result.matrix
+"""
+$TYPEDEF
+
+Storage for the result of a symmetric coloring algorithm with direct decompression.
+
+Similar to [`DefaultColoringResult`](@ref) but contains an additional [`StarSet`](@ref).
+
+# Fields
+
+$TYPEDFIELDS
+
+# See also
+
+- [`AbstractColoringResult`](@ref)
+"""
+struct StarSetColoringResult{partition,M} <:
+       AbstractColoringResult{:symmetric,partition,:direct,M}
+    matrix::M
+    color::Vector{Int}
+    group::Vector{Vector{Int}}
+    star_set::StarSet
+end
+
+function StarSetColoringResult{partition}(
+    matrix::M, color::Vector{Int}, star_set::StarSet
+) where {partition,M}
+    return StarSetColoringResult{partition,M}(
+        matrix, color, group_by_color(color), star_set
+    )
+end
+
+"""
+$TYPEDEF
+
+Storage for the result of a symmetric coloring algorithm with decompression by substitution.
 
-column_colors(result::DefaultColoringResult{s,:column}) where {s} = result.color
-column_groups(result::DefaultColoringResult{s,:column}) where {s} = result.group
+Similar to [`DefaultColoringResult`](@ref) but contains an additional [`TreeSet`](@ref).
 
-row_colors(result::DefaultColoringResult{s,:row}) where {s} = result.color
-row_groups(result::DefaultColoringResult{s,:row}) where {s} = result.group
+# Fields
+
+$TYPEDFIELDS
+
+# See also
+
+- [`AbstractColoringResult`](@ref)
+"""
+struct TreeSetColoringResult{partition,M} <:
+       AbstractColoringResult{:symmetric,partition,:substitution,M}
+    matrix::M
+    color::Vector{Int}
+    group::Vector{Vector{Int}}
+    tree_set::TreeSet
+end
+
+function TreeSetColoringResult{partition}(
+    matrix::M, color::Vector{Int}, tree_set::TreeSet
+) where {partition,M}
+    return TreeSetColoringResult{partition,M}(
+        matrix, color, group_by_color(color), tree_set
+    )
+end

src/sparsematrixcsc.jl

@@ -33,14 +33,6 @@ function DirectSparseColoringResult{structure,partition}(
     )
 end
 
-get_matrix(result::DirectSparseColoringResult) = result.matrix
-
-column_colors(result::DirectSparseColoringResult{s,:column}) where {s} = result.color
-column_groups(result::DirectSparseColoringResult{s,:column}) where {s} = result.group
-
-row_colors(result::DirectSparseColoringResult{s,:row}) where {s} = result.color
-row_groups(result::DirectSparseColoringResult{s,:row}) where {s} = result.group
-
 ## Coloring
 
 function coloring(

test/small.jl

@@ -77,22 +77,24 @@ end;
     @testset "Substitution - Fig 6.1 from 'What color is your Jacobian'" begin
         example = what_fig_61()
         A0, B0, color0 = example.A, example.B, example.color
+        result0 = DefaultColoringResult{:symmetric,:column,:substitution}(A0, color0)
         result = coloring(
             A0,
             ColoringProblem(;
                 structure=:symmetric, partition=:column, decompression=:substitution
             ),
             GreedyColoringAlgorithm(),
-        )
-        color = column_colors(result)
+        )  # returns a TreeSetColoringResult
         group = column_groups(result)
         B = stack(group; dims=2) do g
             dropdims(sum(A0[:, g]; dims=2); dims=2)
         end
-        @test color != color0
+        @test column_colors(result) != color0
         @test B != B0
         @test decompress(B, result) ≈ A0
+        @test decompress(B0, result0) ≈ A0
         for A in matrix_versions(A0)
+            @test decompress!(respectful_similar(A), B0, result0) ≈ A
             @test decompress!(respectful_similar(A), B, result) ≈ A
         end
     end
@@ -112,16 +114,19 @@ end;
     @testset "Substitution - Fig 4 from 'Efficient computation of sparse hessians using coloring and AD'" begin
         example = efficient_fig_4()
         A0, B0, color0 = example.A, example.B, example.color
+        result0 = DefaultColoringResult{:symmetric,:column,:substitution}(A0, color0)
         result = coloring(
             A0,
             ColoringProblem(;
                 structure=:symmetric, partition=:column, decompression=:substitution
             ),
             GreedyColoringAlgorithm(),
-        )
+        )  # returns a TreeSetColoringResult
         @test column_colors(result) == color0
+        @test decompress(B0, result0) ≈ A0
         @test decompress(B0, result) ≈ A0
         for A in matrix_versions(A0)
+            @test decompress!(respectful_similar(A), B0, result0) ≈ A
             @test decompress!(respectful_similar(A), B0, result) ≈ A
         end
     end