BandedBlockBandedMatrices

Author: gdalle

ext/SparseMatrixColoringsBlockBandedMatricesExt.jl

@@ -3,6 +3,7 @@ module SparseMatrixColoringsBlockBandedMatricesExt
 if isdefined(Base, :get_extension)
     using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
     using BlockBandedMatrices:
+        BandedBlockBandedMatrix,
         BlockBandedMatrix,
         blockbandrange,
         blockbandwidths,
@@ -22,6 +23,7 @@ if isdefined(Base, :get_extension)
 else
     using ..BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
     using ..BlockBandedMatrices:
+        BandedBlockBandedMatrix,
         BlockBandedMatrix,
         blockbandrange,
         blockbandwidths,
@@ -46,9 +48,14 @@ https://github.com/JuliaArrays/ArrayInterface.jl
 https://github.com/JuliaDiff/FiniteDiff.jl
 =#
 
-## BlockBandedMatrix
+function subblockbandrange(A::BandedBlockBandedMatrix)
+    u, l = subblockbandwidths(A)
+    return (-l):u
+end
 
-function blockbanded_coloring(A::BlockBandedMatrix, dim::Integer)
+function blockbanded_coloring(
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix}, dim::Integer
+)
     # consider blocks of columns or rows (let's call them vertices) depending on `dim`
     nb_blocks = blocksize(A, dim)
     nb_in_block = blocklengths(axes(A, dim))
@@ -61,27 +68,34 @@ function blockbanded_coloring(A::BlockBandedMatrix, dim::Integer)
     nb_macrocolors = length(blockbandrange(A))
     macrocolor = cycle_range(nb_macrocolors, nb_blocks)
 
+    width = if A isa BandedBlockBandedMatrix
+        # vertices within a block are colored cleverly using bands
+        length(subblockbandrange(A))
+    else
+        # vertices within a block are colored naively with distinct micro colors (~ infinite band width)
+        minimum(size(A))
+    end
+
     # for each macroscopic color, count how many microscopic colors will be needed
-    # vertices within a block are colored naively with all distinct micro colors
-    nb_colors_in_macrocolor = [
-        maximum(nb_in_block[mc:nb_macrocolors:nb_blocks]; init=0) for mc in 1:nb_macrocolors
-    ]
+    nb_colors_in_macrocolor = zeros(Int, nb_macrocolors)
+    for mc in 1:nb_macrocolors
+        largest_nb_in_macrocolor = maximum(nb_in_block[mc:nb_macrocolors:nb_blocks]; init=0)
+        nb_colors_in_macrocolor[mc] = min(width, largest_nb_in_macrocolor)
+    end
     color_shift_in_macrocolor = vcat(0, cumsum(nb_colors_in_macrocolor)[1:(end - 1)])
 
     # assign a microscopic color to each column as a function of its macroscopic color and its position within the block
     for b in 1:nb_blocks
-        mc = macrocolor[b]
-        shift = color_shift_in_macrocolor[mc]
-        for k in first_in_block[b]:last_in_block[b]
-            color[k] = shift + (k - first_in_block[b] + 1)
-        end
+        block_color = cycle_range(width, nb_in_block[b])
+        shift = color_shift_in_macrocolor[macrocolor[b]]
+        color[first_in_block[b]:last_in_block[b]] .= shift .+ block_color
     end
 
     return color
 end
 
 function SMC.coloring(
-    A::BlockBandedMatrix,
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
     ::ColoringProblem{:nonsymmetric,:column},
     algo::GreedyColoringAlgorithm;
     kwargs...,
@@ -92,7 +106,7 @@ function SMC.coloring(
 end
 
 function SMC.coloring(
-    A::BlockBandedMatrix,
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
     ::ColoringProblem{:nonsymmetric,:row},
     algo::GreedyColoringAlgorithm;
     kwargs...,
@@ -102,6 +116,4 @@ function SMC.coloring(
     return RowColoringResult(A, bg, color)
 end
 
-
-
 end

test/structured.jl

@@ -18,14 +18,14 @@ using Test
 end;
 
 @testset "Diagonal" begin
-    @testset for n in (1, 2, 10, 100)
+    for n in (1, 2, 10, 100)
         A = Diagonal(rand(n))
         test_structured_coloring_decompression(A)
     end
 end;
 
 @testset "Bidiagonal" begin
-    @testset for n in (2, 10, 100)
+    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)
@@ -48,10 +48,23 @@ end;
 end;
 
 @testset "BlockBandedMatrices" begin
-    @testset for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, bs in 1:3
-        rows = rand(1:bs, mb)
-        cols = rand(1:bs, nb)
+    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, (l, u))
         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/utils.jl

@@ -127,24 +127,20 @@ function test_structured_coloring_decompression(A::AbstractMatrix)
     row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
     algo = GreedyColoringAlgorithm()
 
-    @testset "Column" begin
-        result = coloring(A, column_problem, algo)
-        color = column_colors(result)
-        B = compress(A, result)
-        D = decompress(B, result)
-        @test D == A
-        @test D isa typeof(A)
-        @test structurally_orthogonal_columns(A, color)
-        if A isa OptimalColoringKnown
-            @test color == ArrayInterface.matrix_colors(A)
-        end
-    end
-
-    @testset "Row" begin
-        result = coloring(A, row_problem, algo)
-        B = compress(A, result)
-        D = decompress(B, result)
-        @test D == A
-        @test D isa typeof(A)
-    end
+    # 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)
+    @test color == ArrayInterface.matrix_colors(A)
+
+    # 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