Add decompression utilities (#4)

* Add decompression utilities * Fix star coloring tests * Fix coverage
gdalle · May 21, 2024 · aaba292 · aaba292
1 parent 35c22cd
commit aaba292
Show file tree

Hide file tree

Showing 13 changed files with 262 additions and 69 deletions.
diff --git a/Project.toml b/Project.toml
@@ -5,13 +5,15 @@ version = "0.1.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 ADTypes = "1.2.1"
+Compat = "3,4"
 DocStringExtensions = "0.9"
 LinearAlgebra = "1"
 Random = "1"

diff --git a/docs/src/api.md b/docs/src/api.md
@@ -2,18 +2,61 @@
 
 ```@meta
 CollapsedDocStrings = true
+CurrentModule = SparseMatrixColorings
 ```
 
-## Public
+## Public, exported
 
-```@autodocs
-Modules = [SparseMatrixColorings]
-Private = false
+```@docs
+SparseMatrixColorings
+GreedyColoringAlgorithm
+```
+
+## Public, not exported
+
+### Orders
+
+```@docs
+AbstractOrder
+NaturalOrder
+RandomOrder
+LargestFirst
+```
+
+### Decompression
+
+```@docs
+decompress_columns!
+decompress_columns
+decompress_rows!
+decompress_rows
 ```
 
 ## Private
 
-```@autodocs
-Modules = [SparseMatrixColorings]
-Public = false
+### Graphs
+
+```@docs
+Graph
+BipartiteGraph
+adjacency_graph
+bipartite_graph
+neighbors
+vertices
+```
+
+### Coloring
+
+```@docs
+partial_distance2_coloring
+star_coloring1
+color_groups
+```
+
+### Testing
+
+```@docs
+check_structurally_orthogonal_columns
+check_structurally_orthogonal_rows
+check_symmetrically_orthogonal
 ```
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
@@ -6,7 +6,8 @@ $README
 module SparseMatrixColorings
 
 using ADTypes: ADTypes, AbstractColoringAlgorithm
-using DocStringExtensions
+using Compat: @compat
+using DocStringExtensions: README
 using LinearAlgebra: Diagonal, Transpose, checksquare, parent, transpose
 using Random: AbstractRNG, default_rng, randperm
 using SparseArrays:
@@ -15,6 +16,7 @@ using SparseArrays:
     dropzeros,
     dropzeros!,
     nnz,
+    nonzeros,
     nzrange,
     rowvals,
     sparse,
@@ -25,6 +27,11 @@ include("order.jl")
 include("coloring.jl")
 include("adtypes.jl")
 include("check.jl")
+include("decompression.jl")
+
+@compat public GreedyColoringAlgorithm
+@compat public NaturalOrder, RandomOrder, LargestFirst
+@compat public decompress_columns!, decompress_columns, decompress_rows!, decompress_rows
 
 export GreedyColoringAlgorithm
 

diff --git a/src/decompression.jl b/src/decompression.jl
@@ -0,0 +1,141 @@
+transpose_respecting_similar(A::AbstractMatrix, ::Type{T}) where {T} = similar(A, T)
+
+function transpose_respecting_similar(A::Transpose, ::Type{T}) where {T}
+    return transpose(similar(parent(A), T))
+end
+
+"""
+    color_groups(colors)
+
+Return `groups::Vector{Vector{Int}}` such that `i ∈ groups[c]` iff `colors[i] == c`.
+
+Assumes the colors are contiguously numbered from `1` to some `cmax`.
+"""
+function color_groups(colors::AbstractVector{<:Integer})
+    cmin, cmax = extrema(colors)
+    @assert cmin == 1
+    groups = [Int[] for c in 1:cmax]
+    for (k, c) in enumerate(colors)
+        push!(groups[c], k)
+    end
+    return groups
+end
+
+## Column decompression
+
+"""
+    decompress_columns!(
+        A::AbstractMatrix{R}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
+    ) where {R<:Real}
+
+Decompress the thin matrix `C` into the fat matrix `A`.
+
+Here, `C` is a compressed representation of matrix `A` obtained by summing the columns that share the same color in `colors`.
+"""
+function decompress_columns! end
+
+#=
+function decompress_columns!(
+    A::AbstractMatrix{R}, 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])
+        copyto!(A[rows_j, j], C[rows_j, k])
+    end
+    return A
+end
+=#
+
+function decompress_columns!(
+    A::SparseMatrixCSC{R}, C::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
+) where {R<:Real}
+    Anz, Arv = nonzeros(A), rowvals(A)
+    Anz .= zero(R)
+    @views for j in axes(A, 2)
+        k = colors[j]
+        nzrange_j = nzrange(A, j)
+        rows_j = Arv[nzrange_j]
+        copyto!(Anz[nzrange_j], C[rows_j, k])
+    end
+    return A
+end
+
+"""
+    decompress_columns(
+        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`.
+"""
+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)
+end
+
+## Row decompression
+
+"""
+    decompress_rows!(
+        A::AbstractMatrix{R},
+        C::AbstractMatrix{R}, S::AbstractMatrix{Bool},
+        colors::AbstractVector{<:Integer}
+    ) where {R<:Real}
+
+Decompress the small matrix `C` into the tall matrix `A`.
+
+Here, `C` is a compressed representation of matrix `A` obtained by summing the rows that share the same color in `colors`.
+"""
+function decompress_rows! end
+
+#=
+function decompress_rows!(
+    A::AbstractMatrix{R}, 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, :])
+        copyto!(A[i, cols_i], C[k, cols_i])
+    end
+    return A
+end
+=#
+
+function decompress_rows!(
+    A::Transpose{R,<:SparseMatrixCSC{R}},
+    C::AbstractMatrix{R},
+    colors::AbstractVector{<:Integer},
+) where {R<:Real}
+    PA = parent(A)
+    PAnz, PArv = nonzeros(PA), rowvals(PA)
+    PAnz .= zero(R)
+    @views for i in axes(A, 1)
+        k = colors[i]
+        nzrange_i = nzrange(PA, i)
+        cols_i = PArv[nzrange_i]
+        copyto!(PAnz[nzrange_i], C[k, cols_i])
+    end
+    return A
+end
+
+"""
+    decompress_rows(
+        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`.
+"""
+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)
+end
diff --git a/src/graph.jl b/src/graph.jl
@@ -110,20 +110,3 @@ function bipartite_graph(J::SparseMatrixCSC)
 end
 
 bipartite_graph(J::AbstractMatrix) = bipartite_graph(sparse(J))
-
-"""
-    column_intersection_graph(J::AbstractMatrix)
-
-Return a [`Graph`](@ref) representing the column intersections of a non-symmetric matrix (typically a Jacobian matrix).
-
-The column intersection graph of a matrix `A ∈ ℝ^{m × n}` is `Gc(A) = (V, E)` where
-
-- `V = 1:n` is the set of columns `j`
-- `(j1, j2) ∈ E` whenever `A[:, j1] ∩ A[:, j2] ≠ ∅`
-"""
-function column_intersection_graph(J::SparseMatrixCSC)
-    A = transpose(J) * J
-    return adjacency_graph(A - Diagonal(A))
-end
-
-column_intersection_graph(J::AbstractMatrix) = column_intersection_graph(sparse(J))
diff --git a/test/Project.toml b/test/Project.toml
@@ -3,6 +3,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"

diff --git a/test/adtypes.jl → test/coloring_correctness.jl b/test/adtypes.jl → test/coloring_correctness.jl
@@ -8,11 +8,13 @@ using SparseMatrixColorings:
     check_symmetrically_orthogonal
 using Test
 
-alg = GreedyColoringAlgorithm()
+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))
-        column_colors = column_coloring(A, alg)
+        column_colors = column_coloring(A, algo)
         @test check_structurally_orthogonal_columns(A, column_colors)
         @test minimum(column_colors) == 1
         @test maximum(column_colors) < size(A, 2) ÷ 2
@@ -21,7 +23,7 @@ end
 
 @testset "Row coloring" begin
     for A in (sprand(Bool, 100, 200, 0.05), sprand(Bool, 200, 100, 0.05))
-        row_colors = row_coloring(A, alg)
+        row_colors = row_coloring(A, algo)
         @test check_structurally_orthogonal_rows(A, row_colors)
         @test minimum(row_colors) == 1
         @test maximum(row_colors) < size(A, 1) ÷ 2
@@ -30,7 +32,7 @@ end
 
 @testset "Symmetric coloring" begin
     S = sparse(Symmetric(sprand(Bool, 100, 100, 0.05)))
-    symmetric_colors = symmetric_coloring(S, alg)
+    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

diff --git a/test/performance.jl → test/coloring_performance.jl b/test/performance.jl → test/coloring_performance.jl
diff --git a/test/decompression_correctness.jl b/test/decompression_correctness.jl
@@ -0,0 +1,38 @@
+using ADTypes: column_coloring, row_coloring, symmetric_coloring
+using Compat
+using SparseArrays
+using SparseMatrixColorings: color_groups, decompress_columns, decompress_rows
+using StableRNGs
+using Test
+
+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)
+
+@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)
+    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)
+    end
+    At_new = decompress_rows(St, Ct, colors)
+    @test At_new == At
+end
diff --git a/test/reference/what_table_41_42.csv b/test/reference/what_table_41_42.csv
@@ -1,7 +1,7 @@
-row,modified,group,name,V,E,∆,δ,K_d2,K_star1,K_star2
+row,modified,group,name,V,E,Δ,δ,K_d2,K_star1,K_star2
 1,false,,mrng1,257000,505048,4,2,12,8,10
 2,false,,mrng2,1017253,2015714,4,2,12,9,10
-3,false,DIMACS10,598a,110971,741934,26,5 13,38,27,32
+3,false,DIMACS10,598a,110971,741934,26,5,38,27,32
 4,false,DIMACS10,144,144649,1074393,26,4,41,28,35
 5,false,DIMACS10,m14b,214765,1679018,40,4,42,29,34
 6,false,DIMACS10,auto,448695,3314611,37,4,42,29,36
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -36,16 +36,23 @@ using Test
         end
     end
     @testset verbose = true "Correctness" begin
-        @testset "ADTypes" begin
-            include("adtypes.jl")
+        @testset "Coloring" begin
+            include("coloring_correctness.jl")
         end
-        @testset "SuiteSparse" begin
-            include("suitesparse.jl")
+        @testset "Decompression" begin
+            include("decompression_correctness.jl")
         end
     end
     @testset "Performance" begin
         if VERSION >= v"1.10"
-            include("performance.jl")
+            @testset "Coloring" begin
+                include("coloring_performance.jl")
+            end
+        end
+    end
+    @testset "Comparison" begin
+        @testset "SuiteSparse" begin
+            include("suitesparse.jl")
         end
     end
 end