Add some functionality to the API (#7)

* Add some functionality to the API

* Fix

Author: gdalle

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseMatrixColorings"
 uuid = "0a514795-09f3-496d-8182-132a7b665d35"
 authors = ["Guillaume Dalle <22795598+gdalle@users.noreply.github.com>"]
-version = "0.2.0"
+version = "0.3.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -15,7 +15,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 ADTypes = "1.2.1"
 Compat = "3,4"
 DocStringExtensions = "0.9"
-LinearAlgebra = "1"
-Random = "1"
-SparseArrays = "1"
+LinearAlgebra = "<0.0.1, 1"
+Random = "<0.0.1, 1"
+SparseArrays = "<0.0.1, 1"
 julia = "1.6"

docs/src/api.md

@@ -30,6 +30,7 @@ decompress_columns!
 decompress_columns
 decompress_rows!
 decompress_rows
+color_groups
 ```
 
 ## Private
@@ -50,7 +51,6 @@ vertices
 ```@docs
 partial_distance2_coloring
 star_coloring1
-color_groups
 ```
 
 ### Testing

src/SparseMatrixColorings.jl

@@ -31,7 +31,9 @@ include("decompression.jl")
 
 @compat public GreedyColoringAlgorithm
 @compat public NaturalOrder, RandomOrder, LargestFirst
-@compat public decompress_columns!, decompress_columns, decompress_rows!, decompress_rows
+@compat public decompress_columns!, decompress_columns
+@compat public decompress_rows!, decompress_rows
+@compat public color_groups
 
 export GreedyColoringAlgorithm
 

src/decompression.jl

@@ -25,32 +25,41 @@ end
 
 """
     decompress_columns!(
-        A::AbstractMatrix{R}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
+        A::AbstractMatrix{R},
+        S::AbstractMatrix{Bool},
+        C::AbstractMatrix{R},
+        colors::AbstractVector{<:Integer}
     ) where {R<:Real}
 
-Decompress the thin matrix `C` into the fat matrix `A`.
+Decompress the thin matrix `C` into a fat matrix `A` with the same sparsity pattern as `S`.
 
-Here, `C` is a compressed representation of matrix `A` obtained by summing the columns that share the same color in `colors`.
+Here, `colors` is a column coloring of `S`, while `C` is a compressed representation of matrix `A` obtained by summing the columns that share the same color.
 """
 function decompress_columns! end
 
-#=
 function decompress_columns!(
-    A::AbstractMatrix{R}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
+    A::AbstractMatrix{R},
+    S::AbstractMatrix{Bool},
+    C::AbstractMatrix{R},
+    colors::AbstractVector{<:Integer},
 ) where {R<:Real}
     A .= zero(R)
     @views for j in axes(A, 2)
         k = colors[j]
-        rows_j = (!iszero).(A[:, j])
+        rows_j = map(!iszero, S[:, j])
         copyto!(A[rows_j, j], C[rows_j, k])
+        A[rows_j, j] .= C[rows_j, k]
     end
     return A
 end
-=#
 
 function decompress_columns!(
-    A::SparseMatrixCSC{R}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
+    A::SparseMatrixCSC{R},
+    S::SparseMatrixCSC{Bool},
+    C::AbstractMatrix{R},
+    colors::AbstractVector{<:Integer},
 ) where {R<:Real}
+    # assume A and S have the same pattern
     Anz, Arv = nonzeros(A), rowvals(A)
     Anz .= zero(R)
     @views for j in axes(A, 2)
@@ -64,54 +73,60 @@ end
 
 """
     decompress_columns(
-        S::AbstractMatrix{Bool}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
+        S::AbstractMatrix{Bool},
+        C::AbstractMatrix{R},
+        colors::AbstractVector{<:Integer}
     ) where {R<:Real}
 
 Decompress the thin matrix `C` into a new fat matrix `A` with the same sparsity pattern as `S`.
 
-Here, `C` is a compressed representation of matrix `A` obtained by summing the columns that share the same color in `colors`.
+Here, `colors` is a column coloring of `S`, while `C` is a compressed representation of matrix `A` obtained by summing the columns that share the same color.
 """
 function decompress_columns(
     S::AbstractMatrix{Bool}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
 ) where {R<:Real}
     A = transpose_respecting_similar(S, R)
-    return decompress_columns!(A, C, colors)
+    return decompress_columns!(A, S, C, colors)
 end
 
 ## Row decompression
 
 """
     decompress_rows!(
         A::AbstractMatrix{R},
-        C::AbstractMatrix{R}, S::AbstractMatrix{Bool},
+        S::AbstractMatrix{Bool},
+        C::AbstractMatrix{R},
         colors::AbstractVector{<:Integer}
     ) where {R<:Real}
 
-Decompress the small matrix `C` into the tall matrix `A`.
+Decompress the small matrix `C` into the tall matrix `A` with the same sparsity pattern as `S`.
 
-Here, `C` is a compressed representation of matrix `A` obtained by summing the rows that share the same color in `colors`.
+Here, `colors` is a row coloring of `S`, while `C` is a compressed representation of matrix `A` obtained by summing the columns that share the same color.
 """
 function decompress_rows! end
 
-#=
 function decompress_rows!(
-    A::AbstractMatrix{R}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
+    A::AbstractMatrix{R},
+    S::AbstractMatrix{Bool},
+    C::AbstractMatrix{R},
+    colors::AbstractVector{<:Integer},
 ) where {R<:Real}
     A .= zero(R)
     @views for i in axes(A, 1)
         k = colors[i]
-        cols_i = (!iszero).(A[i, :])
+        cols_i = map(!iszero, S[i, :])
         copyto!(A[i, cols_i], C[k, cols_i])
     end
     return A
 end
-=#
 
 function decompress_rows!(
     A::Transpose{R,<:SparseMatrixCSC{R}},
+    S::Transpose{Bool,<:SparseMatrixCSC{Bool}},
     C::AbstractMatrix{R},
     colors::AbstractVector{<:Integer},
 ) where {R<:Real}
+    # assume A and S have the same pattern
     PA = parent(A)
     PAnz, PArv = nonzeros(PA), rowvals(PA)
     PAnz .= zero(R)
@@ -126,16 +141,18 @@ end
 
 """
     decompress_rows(
-        S::AbstractMatrix{Bool}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
+        S::AbstractMatrix{Bool},
+        C::AbstractMatrix{R},
+        colors::AbstractVector{<:Integer}
     ) where {R<:Real}
 
 Decompress the small matrix `C` into a new tall matrix `A` with the same sparsity pattern as `S`.
 
-Here, `C` is a compressed representation of matrix `A` obtained by summing the rows that share the same color in `colors`.
+Here, `colors` is a row coloring of `S`, while `C` is a compressed representation of matrix `A` obtained by summing the columns that share the same color.
 """
 function decompress_rows(
     S::AbstractMatrix{Bool}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
 ) where {R<:Real}
     A = transpose_respecting_similar(S, R)
-    return decompress_rows!(A, C, colors)
+    return decompress_rows!(A, S, C, colors)
 end

test/coloring_correctness.jl

@@ -6,14 +6,22 @@ using SparseMatrixColorings:
     check_structurally_orthogonal_columns,
     check_structurally_orthogonal_rows,
     check_symmetrically_orthogonal
+using StableRNGs
 using Test
 
+rng = StableRNG(63)
+
 algo = GreedyColoringAlgorithm()
 
 @test startswith(string(algo), "GreedyColoringAlgorithm(")
 
 @testset "Column coloring" begin
-    for A in (sprand(Bool, 100, 200, 0.05), sprand(Bool, 200, 100, 0.05))
+    @testset "$(typeof(A)) - $(size(A))" for A in (
+        sprand(rng, Bool, 100, 200, 0.05),
+        sprand(rng, Bool, 200, 100, 0.05),
+        Matrix(sprand(rng, Bool, 100, 200, 0.05)),
+        Matrix(sprand(rng, Bool, 200, 100, 0.05)),
+    )
         column_colors = column_coloring(A, algo)
         @test check_structurally_orthogonal_columns(A, column_colors)
         @test minimum(column_colors) == 1
@@ -22,7 +30,12 @@ algo = GreedyColoringAlgorithm()
 end
 
 @testset "Row coloring" begin
-    for A in (sprand(Bool, 100, 200, 0.05), sprand(Bool, 200, 100, 0.05))
+    @testset "$(typeof(A)) - $(size(A))" for A in (
+        sprand(rng, Bool, 100, 200, 0.05),
+        sprand(rng, Bool, 200, 100, 0.05),
+        Matrix(sprand(rng, Bool, 100, 200, 0.05)),
+        Matrix(sprand(rng, Bool, 200, 100, 0.05)),
+    )
         row_colors = row_coloring(A, algo)
         @test check_structurally_orthogonal_rows(A, row_colors)
         @test minimum(row_colors) == 1
@@ -31,9 +44,13 @@ end
 end
 
 @testset "Symmetric coloring" begin
-    S = sparse(Symmetric(sprand(Bool, 100, 100, 0.05)))
-    symmetric_colors = symmetric_coloring(S, algo)
-    @test check_symmetrically_orthogonal(S, symmetric_colors)
-    @test minimum(symmetric_colors) == 1
-    @test maximum(symmetric_colors) < size(S, 2) ÷ 2
+    @testset "$(typeof(A)) - $(size(A))" for A in (
+        sparse(Symmetric(sprand(rng, Bool, 100, 100, 0.05))),
+        Symmetric(Matrix(sprand(rng, Bool, 100, 100, 0.05))),
+    )
+        symmetric_colors = symmetric_coloring(A, algo)
+        @test check_symmetrically_orthogonal(A, symmetric_colors)
+        @test minimum(symmetric_colors) == 1
+        @test maximum(symmetric_colors) < size(A, 2) ÷ 2
+    end
 end

test/decompression_correctness.jl

@@ -1,6 +1,7 @@
 using ADTypes: column_coloring, row_coloring, symmetric_coloring
 using Compat
 using SparseArrays
+using SparseMatrixColorings
 using SparseMatrixColorings: color_groups, decompress_columns, decompress_rows
 using StableRNGs
 using Test
@@ -10,29 +11,39 @@ rng = StableRNG(63)
 algo = GreedyColoringAlgorithm()
 
 m, n = 10, 20
-A = sprand(rng, Bool, m, n, 0.3)
-At = transpose(A)
-S = map(!iszero, A)
-St = transpose(S)
+
+A0 = sprand(rng, Bool, m, n, 0.3)
+A1 = Matrix(A0)
+A0t = transpose(A0)
+A1t = transpose(A1)
+
+S0 = map(!iszero, A0)
+S1 = map(!iszero, A1)
+S0t = transpose(S0)
+S1t = transpose(S1)
 
 @testset "Column decompression" begin
-    colors = column_coloring(A, algo)
-    groups = color_groups(colors)
-    @test length(groups[1]) > 1
-    C = stack(groups) do group
-        dropdims(sum(A[:, group]; dims=2); dims=2)
+    @testset "$(typeof(A))" for (A, S) in zip((A0, A1), (S0, S1))
+        colors = column_coloring(A, algo)
+        groups = color_groups(colors)
+        @test length(groups[1]) > 1
+        C = stack(groups) do group
+            dropdims(sum(A[:, group]; dims=2); dims=2)
+        end
+        A_new = decompress_columns(S, C, colors)
+        @test A_new == A
     end
-    A_new = decompress_columns(S, C, colors)
-    @test A_new == A
 end
 
 @testset "Row decompression" begin
-    colors = row_coloring(At, algo)
-    groups = color_groups(colors)
-    @test length(groups[1]) > 1
-    Ct = stack(groups; dims=1) do group
-        dropdims(sum(At[group, :]; dims=1); dims=1)
+    @testset "$(typeof(At))" for (At, St) in zip((A0t, A1t), (S0t, S1t))
+        colors = row_coloring(At, algo)
+        groups = color_groups(colors)
+        @test length(groups[1]) > 1
+        Ct = stack(groups; dims=1) do group
+            dropdims(sum(At[group, :]; dims=1); dims=1)
+        end
+        At_new = decompress_rows(St, Ct, colors)
+        @test At_new == At
     end
-    At_new = decompress_rows(St, Ct, colors)
-    @test At_new == At
 end