Better sparsity pattern check (#9)

* Better sparsity pattern check

* Coverage

Author: gdalle

src/decompression.jl

@@ -4,6 +4,30 @@ function transpose_respecting_similar(A::Transpose, ::Type{T}) where {T}
     return transpose(similar(parent(A), T))
 end
 
+function same_sparsity_pattern(A::SparseMatrixCSC, B::SparseMatrixCSC)
+    if size(A) != size(B)
+        return false
+    elseif nnz(A) != nnz(B)
+        return false
+    else
+        for j in axes(A, 2)
+            rA = nzrange(A, j)
+            rB = nzrange(B, j)
+            if rA != rB
+                return false
+            end
+            # TODO: check rowvals?
+        end
+        return true
+    end
+end
+
+function same_sparsity_pattern(
+    A::Transpose{<:Any,<:SparseMatrixCSC}, B::Transpose{<:Any,<:SparseMatrixCSC}
+)
+    return same_sparsity_pattern(parent(A), parent(B))
+end
+
 """
     color_groups(colors)
 
@@ -60,7 +84,7 @@ function decompress_columns!(
     C::AbstractMatrix{R},
     colors::AbstractVector{<:Integer},
 ) where {R<:Real}
-    if nnz(parent(A)) != nnz(parent(S))
+    if !same_sparsity_pattern(A, S)
         throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
     end
     Anz, Arv = nonzeros(A), rowvals(A)
@@ -133,7 +157,7 @@ function decompress_rows!(
     C::AbstractMatrix{R},
     colors::AbstractVector{<:Integer},
 ) where {R<:Real}
-    if nnz(parent(A)) != nnz(parent(S))
+    if !same_sparsity_pattern(A, S)
         throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
     end
     PA = parent(A)

test/decompression_correctness.jl

@@ -1,15 +1,55 @@
 using ADTypes: column_coloring, row_coloring, symmetric_coloring
 using Compat
+using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
-using SparseMatrixColorings: color_groups, decompress_columns, decompress_rows
+using SparseMatrixColorings:
+    color_groups,
+    decompress_columns,
+    decompress_columns!,
+    decompress_rows,
+    decompress_rows!,
+    same_sparsity_pattern
 using StableRNGs
 using Test
 
 rng = StableRNG(63)
 
 algo = GreedyColoringAlgorithm()
 
+@testset "Sparsity pattern comparison" begin
+    A = [
+        1 1
+        0 1
+        0 0
+    ]
+    B1 = [
+        1 1
+        0 1
+        1 0
+    ]
+    B2 = [
+        1 1
+        0 0
+        0 1
+    ]
+    C = [
+        1 1 0
+        0 1 0
+        0 0 0
+    ]
+
+    @test same_sparsity_pattern(sparse(A), sparse(A))
+    @test !same_sparsity_pattern(sparse(A), sparse(B1))
+    @test_broken !same_sparsity_pattern(sparse(A), sparse(B2))
+    @test !same_sparsity_pattern(sparse(A), sparse(C))
+
+    @test same_sparsity_pattern(transpose(sparse(A)), transpose(sparse(A)))
+    @test !same_sparsity_pattern(transpose(sparse(A)), transpose(sparse(B1)))
+    @test_broken !same_sparsity_pattern(transpose(sparse(A)), transpose(sparse(B2)))
+    @test !same_sparsity_pattern(transpose(sparse(A)), transpose(sparse(C)))
+end;
+
 @testset "Column decompression" begin
     @testset "Small" begin
         A0 = [
@@ -33,6 +73,11 @@ algo = GreedyColoringAlgorithm()
             (sparse(A0), sparse(S0)),
         ]
             @test decompress_columns(S, C, colors) == A
+            if A isa SparseMatrixCSC
+                @test_throws DimensionMismatch decompress_columns!(
+                    similar(A), false .* S, C, colors
+                )
+            end
         end
     end
     @testset "Medium" begin
@@ -76,6 +121,11 @@ end
             (transpose(sparse(transpose(A0))), transpose(sparse(transpose(S0)))),
         ]
             @test decompress_rows(S, C, colors) == A
+            if A isa Transpose{<:Any,<:SparseMatrixCSC}
+                @test_throws DimensionMismatch decompress_rows!(
+                    transpose(similar(parent(A))), transpose(false .* parent(S)), C, colors
+                )
+            end
         end
     end
     @testset "Medium" begin