Coloring and decompression for structured matrices

src/check.jl

@@ -36,7 +36,7 @@ function structurally_orthogonal_columns(
     group = group_by_color(color)
     for (c, g) in enumerate(group)
         Ag = view(A, :, g)
-        nonzeros_per_row = dropdims(count(!iszero, Ag; dims=2); dims=2)
+        nonzeros_per_row = only(eachcol(count(!iszero, Ag; dims=2)))
         max_nonzeros_per_row, i = findmax(nonzeros_per_row)
         if max_nonzeros_per_row > 1
             if verbose

test/Project.toml

@@ -1,6 +1,10 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

test/allocations.jl

@@ -2,7 +2,7 @@ using Chairmarks
 using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
-using SparseMatrixColorings: partial_distance2_coloring!
+using SparseMatrixColorings: BipartiteGraph, partial_distance2_coloring!
 using StableRNGs
 using Test
 
@@ -21,7 +21,7 @@ end
     test_noallocs_distance2_coloring(1000)
 end;
 
-function test_noallocs_decompression(
+function test_noallocs_sparse_decompression(
     n::Integer; structure::Symbol, partition::Symbol, decompression::Symbol
 )
     A = sparse(Symmetric(sprand(rng, n, n, 5 / n)))
@@ -75,7 +75,27 @@ function test_noallocs_decompression(
     end
 end
 
-@testset "Decompression" begin
+function test_noallocs_structured_decompression(
+    n::Integer; structure::Symbol, partition::Symbol, decompression::Symbol
+)
+    @testset "$(nameof(typeof(A)))" for A in [
+        Diagonal(rand(n)),
+        Bidiagonal(rand(n), rand(n - 1), 'U'),
+        Bidiagonal(rand(n), rand(n - 1), 'L'),
+        Tridiagonal(rand(n - 1), rand(n), rand(n - 1)),
+    ]
+        result = coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression),
+        )
+        B = compress(A, result)
+        bench = @be similar(A) decompress!(_, B, result) evals = 1
+        @test minimum(bench).allocs == 0
+    end
+end
+
+@testset "Sparse decompression" begin
     @testset "$structure - $partition - $decompression" for (
         structure, partition, decompression
     ) in [
@@ -84,6 +104,16 @@ end
         (:symmetric, :column, :direct),
         (:symmetric, :column, :substitution),
     ]
-        test_noallocs_decompression(1000; structure, partition, decompression)
+        test_noallocs_sparse_decompression(1000; structure, partition, decompression)
+    end
+end;
+
+@testset "Structured decompression" begin
+    @testset "$structure - $partition - $decompression" for (
+        structure, partition, decompression
+    ) in [
+        (:nonsymmetric, :column, :direct), (:nonsymmetric, :row, :direct)
+    ]
+        test_noallocs_structured_decompression(1000; structure, partition, decompression)
     end
 end;

test/runtests.jl

@@ -53,6 +53,9 @@ include("utils.jl")
         @testset "Random instances" begin
             include("random.jl")
         end
+        @testset "Structured matrices" begin
+            include("structured.jl")
+        end
         @testset "Instances with known colorings" begin
             include("theory.jl")
         end

test/structured.jl

@@ -0,0 +1,58 @@
+using ArrayInterface: ArrayInterface
+using BandedMatrices: BandedMatrix, brand
+using BlockBandedMatrices: BandedBlockBandedMatrix, BlockBandedMatrix
+using LinearAlgebra
+using SparseMatrixColorings
+using Test
+
+@testset "Diagonal" begin
+    for n in (1, 2, 10, 100)
+        A = Diagonal(rand(n))
+        test_structured_coloring_decompression(A)
+    end
+end;
+
+@testset "Bidiagonal" begin
+    for n in (2, 10, 100)
+        A1 = Bidiagonal(rand(n), rand(n - 1), :U)
+        A2 = Bidiagonal(rand(n), rand(n - 1), :L)
+        test_structured_coloring_decompression(A1)
+        test_structured_coloring_decompression(A2)
+    end
+end;
+
+@testset "Tridiagonal" begin
+    for n in (2, 10, 100)
+        A = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
+        test_structured_coloring_decompression(A)
+    end
+end;
+
+@testset "BandedMatrices" begin
+    @testset for (m, n) in [(10, 20), (20, 10)], l in 0:5, u in 0:5
+        A = brand(m, n, l, u)
+        test_structured_coloring_decompression(A)
+    end
+end;
+
+@testset "BlockBandedMatrices" begin
+    for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, _ in 1:10
+        rows = rand(1:5, mb)
+        cols = rand(1:5, nb)
+        A = BlockBandedMatrix{Float64}(rand(sum(rows), sum(cols)), rows, cols, (lb, ub))
+        test_structured_coloring_decompression(A)
+    end
+end;
+
+@testset "BandedBlockBandedMatrices" begin
+    for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, _ in 1:10
+        rows = rand(5:10, mb)
+        cols = rand(5:10, nb)
+        λ = rand(0:5)
+        μ = rand(0:5)
+        A = BandedBlockBandedMatrix{Float64}(
+            rand(sum(rows), sum(cols)), rows, cols, (lb, ub), (λ, μ)
+        )
+        test_structured_coloring_decompression(A)
+    end
+end;

test/type_stability.jl

@@ -3,13 +3,13 @@ using JET
 using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
-using SparseMatrixColorings: respectful_similar
+using SparseMatrixColorings: matrix_versions, respectful_similar
 using StableRNGs
 using Test
 
 rng = StableRNG(63)
 
-@testset "Coloring" begin
+@testset "Sparse coloring" begin
     n = 10
     A = sprand(rng, n, n, 5 / n)
 
@@ -40,9 +40,28 @@ rng = StableRNG(63)
             GreedyColoringAlgorithm(; decompression),
         )
     end
-end
+end;
 
-@testset "Decompression" begin
+@testset "Structured coloring" begin
+    n = 10
+    @testset "$(nameof(typeof(A))) - $structure - $partition - $decompression" for A in [
+            Diagonal(rand(n)),
+            Bidiagonal(rand(n), rand(n - 1), 'U'),
+            Bidiagonal(rand(n), rand(n - 1), 'L'),
+            Tridiagonal(rand(n - 1), rand(n), rand(n - 1)),
+        ],
+        (structure, partition, decompression) in
+        [(:nonsymmetric, :column, :direct), (:nonsymmetric, :row, :direct)]
+
+        @test_opt target_modules = (SparseMatrixColorings,) coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression),
+        )
+    end
+end;
+
+@testset "Sparse decompression" begin
     n = 10
     A0 = sparse(Symmetric(sprand(rng, n, n, 5 / n)))
 
@@ -92,3 +111,24 @@ end
         end
     end
 end;
+
+@testset "Structured decompression" begin
+    n = 10
+    @testset "$(nameof(typeof(A))) - $structure - $partition - $decompression" for A in [
+            Diagonal(rand(n)),
+            Bidiagonal(rand(n), rand(n - 1), 'U'),
+            Bidiagonal(rand(n), rand(n - 1), 'L'),
+            Tridiagonal(rand(n - 1), rand(n), rand(n - 1)),
+        ],
+        (structure, partition, decompression) in
+        [(:nonsymmetric, :column, :direct), (:nonsymmetric, :row, :direct)]
+
+        result = coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression);
+        )
+        B = compress(A, result)
+        @test_opt decompress(B, result)
+    end
+end;

test/utils.jl

@@ -1,3 +1,6 @@
+using ArrayInterface: ArrayInterface
+using BandedMatrices: BandedMatrix
+using BlockBandedMatrices: BlockBandedMatrix
 using LinearAlgebra
 using SparseMatrixColorings
 using SparseMatrixColorings:
@@ -138,3 +141,30 @@ function test_coloring_decompression(
         @test all(color_vec .== Ref(color_vec[1]))
     end
 end
+
+function test_structured_coloring_decompression(A::AbstractMatrix)
+    column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+    row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
+    algo = GreedyColoringAlgorithm()
+
+    # Column
+    result = coloring(A, column_problem, algo)
+    color = column_colors(result)
+    B = compress(A, result)
+    D = decompress(B, result)
+    @test D == A
+    @test nameof(typeof(D)) == nameof(typeof(A))
+    @test structurally_orthogonal_columns(A, color)
+    if VERSION >= v"1.10" || A isa Union{Diagonal,Bidiagonal,Tridiagonal}
+        # banded matrices not supported by ArrayInterface on Julia 1.6
+        # @test color == ArrayInterface.matrix_colors(A)  # TODO: uncomment
+        nothing
+    end
+
+    # Row
+    result = coloring(A, row_problem, algo)
+    B = compress(A, result)
+    D = decompress(B, result)
+    @test D == A
+    @test nameof(typeof(D)) == nameof(typeof(A))
+end