JuliaDiff · gdalle · May 24, 2024 · May 24, 2024
diff --git a/Project.toml b/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.3.0"
+version = "0.3.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/src/decompression.jl b/src/decompression.jl
@@ -31,7 +31,7 @@ end
         colors::AbstractVector{<:Integer}
     ) where {R<:Real}
 
-Decompress the thin matrix `C` into a fat matrix `A` with the same sparsity pattern as `S`.
+Decompress the thin matrix `C` into the fat matrix `A` which must have the same sparsity pattern as `S`.
 
 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.
 """
@@ -44,11 +44,12 @@ function decompress_columns!(
     colors::AbstractVector{<:Integer},
 ) where {R<:Real}
     A .= zero(R)
-    @views for j in axes(A, 2)
+    for j in axes(A, 2)
         k = colors[j]
-        rows_j = map(!iszero, S[:, j])
-        copyto!(A[rows_j, j], C[rows_j, k])
-        A[rows_j, j] .= C[rows_j, k]
+        rows_j = (!iszero).(view(S, :, j))
+        Aj = view(A, rows_j, j)
+        Cj = view(C, rows_j, k)
+        copyto!(Aj, Cj)
     end
     return A
 end
@@ -59,14 +60,18 @@ function decompress_columns!(
     C::AbstractMatrix{R},
     colors::AbstractVector{<:Integer},
 ) where {R<:Real}
-    # assume A and S have the same pattern
+    if nnz(parent(A)) != nnz(parent(S))
+        throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
+    end
     Anz, Arv = nonzeros(A), rowvals(A)
     Anz .= zero(R)
-    @views for j in axes(A, 2)
+    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])
+        rows_j = view(Arv, nzrange_j)
+        Aj = view(Anz, nzrange_j)
+        Cj = view(C, rows_j, k)
+        copyto!(Aj, Cj)
     end
     return A
 end
@@ -99,7 +104,7 @@ end
         colors::AbstractVector{<:Integer}
     ) where {R<:Real}
 
-Decompress the small matrix `C` into the tall matrix `A` with the same sparsity pattern as `S`.
+Decompress the small matrix `C` into the tall matrix `A` which must have the same sparsity pattern as `S`.
 
 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.
 """
@@ -112,10 +117,12 @@ function decompress_rows!(
     colors::AbstractVector{<:Integer},
 ) where {R<:Real}
     A .= zero(R)
-    @views for i in axes(A, 1)
+    for i in axes(A, 1)
         k = colors[i]
-        cols_i = map(!iszero, S[i, :])
-        copyto!(A[i, cols_i], C[k, cols_i])
+        cols_i = (!iszero).(view(S, i, :))
+        Ai = view(A, i, cols_i)
+        Ci = view(C, k, cols_i)
+        copyto!(Ai, Ci)
     end
     return A
 end
@@ -126,15 +133,19 @@ function decompress_rows!(
     C::AbstractMatrix{R},
     colors::AbstractVector{<:Integer},
 ) where {R<:Real}
-    # assume A and S have the same pattern
+    if nnz(parent(A)) != nnz(parent(S))
+        throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
+    end
     PA = parent(A)
     PAnz, PArv = nonzeros(PA), rowvals(PA)
     PAnz .= zero(R)
-    @views for i in axes(A, 1)
+    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])
+        cols_i = view(PArv, nzrange_i)
+        Ai = view(PAnz, nzrange_i)
+        Ci = view(C, k, cols_i)
+        copyto!(Ai, Ci)
     end
     return A
 end

diff --git a/test/decompression_correctness.jl b/test/decompression_correctness.jl
@@ -10,40 +10,91 @@ rng = StableRNG(63)
 
 algo = GreedyColoringAlgorithm()
 
-m, n = 10, 20
-
-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
-    @testset "$(typeof(A))" for (A, S) in zip((A0, A1), (S0, S1))
-        colors = column_coloring(A, algo)
+    @testset "Small" begin
+        A0 = [
+            1 0 2
+            0 3 4
+            5 0 0
+        ]
+        S0 = Bool[
+            1 0 1
+            0 1 1
+            1 0 0
+        ]
+        C = [
+            1 2
+            3 4
+            5 0
+        ]
+        colors = [1, 1, 2]
+        @testset "$(typeof(A)) - $(typeof(S))" for (A, S) in [
+            (A0, S0),  #
+            (sparse(A0), sparse(S0)),
+        ]
+            @test decompress_columns(S, C, colors) == A
+        end
+    end
+    @testset "Medium" begin
+        m, n = 18, 20
+        A0 = sprand(rng, Bool, m, n, 0.2)
+        S0 = map(!iszero, A0)
+        colors = column_coloring(A0, algo)
         groups = color_groups(colors)
-        @test length(groups[1]) > 1
-        C = stack(groups) do group
-            dropdims(sum(A[:, group]; dims=2); dims=2)
+        C = stack(groups; dims=2) do group
+            dropdims(sum(A0[:, group]; dims=2); dims=2)
+        end
+        @test size(C) == (size(A0, 1), length(groups))
+        @testset "$(typeof(A)) - $(typeof(S))" for (A, S) in [
+            (A0, S0),  #
+            (Matrix(A0), Matrix(S0)),
+        ]
+            @test decompress_columns(S, C, colors) == A
         end
-        A_new = decompress_columns(S, C, colors)
-        @test A_new == A
     end
 end
 
 @testset "Row decompression" begin
-    @testset "$(typeof(At))" for (At, St) in zip((A0t, A1t), (S0t, S1t))
-        colors = row_coloring(At, algo)
+    @testset "Small" begin
+        A0 = [
+            1 0 3
+            0 2 0
+            4 5 0
+        ]
+        S0 = Bool[
+            1 0 1
+            0 1 0
+            1 1 0
+        ]
+        C = [
+            1 2 3
+            4 5 0
+        ]
+        colors = [1, 1, 2]
+        @testset "$(typeof(A)) - $(typeof(S))" for (A, S) in [
+            (A0, S0), #
+            (transpose(sparse(transpose(A0))), transpose(sparse(transpose(S0)))),
+        ]
+            @test decompress_rows(S, C, colors) == A
+        end
+    end
+    @testset "Medium" begin
+        m, n = 18, 20
+        A0 = sprand(rng, Bool, m, n, 0.2)
+        S0 = map(!iszero, A0)
+        A0t = transpose(A0)
+        S0t = transpose(S0)
+        colors = row_coloring(A0t, 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)
+            dropdims(sum(A0t[group, :]; dims=1); dims=1)
+        end
+        @test size(Ct) == (length(groups), size(A0t, 2))
+        @testset "$(typeof(At)) - $(typeof(St))" for (At, St) in [
+            (A0t, S0t),  #
+            (Matrix(A0t), Matrix(S0t)),
+        ]
+            @test decompress_rows(St, Ct, colors) == At
         end
-        At_new = decompress_rows(St, Ct, colors)
-        @test At_new == At
     end
 end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -43,14 +43,14 @@ using Test
             include("decompression_correctness.jl")
         end
     end
-    @testset "Performance" begin
+    @testset verbose = true "Performance" begin
         if VERSION >= v"1.10"
             @testset "Coloring" begin
                 include("coloring_performance.jl")
             end
         end
     end
-    @testset "Comparison" begin
+    @testset verbose = true "Comparison" begin
         @testset "SuiteSparse" begin
             include("suitesparse.jl")
         end