|
| 1 | +module SparseMatrixColoringsBlockBandedMatricesExt |
| 2 | + |
| 3 | +using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths |
| 4 | +using BlockBandedMatrices: |
| 5 | + BandedBlockBandedMatrix, |
| 6 | + BlockBandedMatrix, |
| 7 | + blockbandrange, |
| 8 | + blockbandwidths, |
| 9 | + blocklengths, |
| 10 | + blocksize, |
| 11 | + subblockbandwidths |
| 12 | +using SparseMatrixColorings: |
| 13 | + BipartiteGraph, |
| 14 | + ColoringProblem, |
| 15 | + ColumnColoringResult, |
| 16 | + StructuredColoringAlgorithm, |
| 17 | + RowColoringResult, |
| 18 | + column_colors, |
| 19 | + cycle_range, |
| 20 | + row_colors |
| 21 | +import SparseMatrixColorings as SMC |
| 22 | + |
| 23 | +#= |
| 24 | +This code is partly taken from ArrayInterface.jl and FiniteDiff.jl |
| 25 | +https://github.com/JuliaArrays/ArrayInterface.jl |
| 26 | +https://github.com/JuliaDiff/FiniteDiff.jl |
| 27 | +=# |
| 28 | + |
| 29 | +function subblockbandrange(A::BandedBlockBandedMatrix) |
| 30 | + u, l = subblockbandwidths(A) |
| 31 | + return (-l):u |
| 32 | +end |
| 33 | + |
| 34 | +function blockbanded_coloring( |
| 35 | + A::Union{BlockBandedMatrix,BandedBlockBandedMatrix}, dim::Integer |
| 36 | +) |
| 37 | + # consider blocks of columns or rows (let's call them vertices) depending on `dim` |
| 38 | + nb_blocks = blocksize(A, dim) |
| 39 | + nb_in_block = blocklengths(axes(A, dim)) |
| 40 | + first_in_block = blockfirsts(axes(A, dim)) |
| 41 | + last_in_block = blocklasts(axes(A, dim)) |
| 42 | + color = zeros(Int, size(A, dim)) |
| 43 | + |
| 44 | + # give a macroscopic color to each block, so that 2 blocks with the same macro color are orthogonal |
| 45 | + # same idea as for BandedMatrices |
| 46 | + nb_macrocolors = length(blockbandrange(A)) |
| 47 | + macrocolor = cycle_range(nb_macrocolors, nb_blocks) |
| 48 | + |
| 49 | + width = if A isa BandedBlockBandedMatrix |
| 50 | + # vertices within a block are colored cleverly using bands |
| 51 | + length(subblockbandrange(A)) |
| 52 | + else |
| 53 | + # vertices within a block are colored naively with distinct micro colors (~ infinite band width) |
| 54 | + typemax(Int) |
| 55 | + end |
| 56 | + |
| 57 | + # for each macroscopic color, count how many microscopic colors will be needed |
| 58 | + nb_colors_in_macrocolor = zeros(Int, nb_macrocolors) |
| 59 | + for mc in 1:nb_macrocolors |
| 60 | + largest_nb_in_macrocolor = maximum(nb_in_block[mc:nb_macrocolors:nb_blocks]; init=0) |
| 61 | + nb_colors_in_macrocolor[mc] = min(width, largest_nb_in_macrocolor) |
| 62 | + end |
| 63 | + color_shift_in_macrocolor = vcat(0, cumsum(nb_colors_in_macrocolor)[1:(end - 1)]) |
| 64 | + |
| 65 | + # assign a microscopic color to each column as a function of its macroscopic color and its position within the block |
| 66 | + for b in 1:nb_blocks |
| 67 | + block_color = cycle_range(width, nb_in_block[b]) |
| 68 | + shift = color_shift_in_macrocolor[macrocolor[b]] |
| 69 | + color[first_in_block[b]:last_in_block[b]] .= shift .+ block_color |
| 70 | + end |
| 71 | + |
| 72 | + return color |
| 73 | +end |
| 74 | + |
| 75 | +function SMC.coloring( |
| 76 | + A::Union{BlockBandedMatrix,BandedBlockBandedMatrix}, |
| 77 | + ::ColoringProblem{:nonsymmetric,:column}, |
| 78 | + ::StructuredColoringAlgorithm; |
| 79 | + kwargs..., |
| 80 | +) |
| 81 | + color = blockbanded_coloring(A, 2) |
| 82 | + bg = BipartiteGraph(A) |
| 83 | + return ColumnColoringResult(A, bg, color) |
| 84 | +end |
| 85 | + |
| 86 | +function SMC.coloring( |
| 87 | + A::Union{BlockBandedMatrix,BandedBlockBandedMatrix}, |
| 88 | + ::ColoringProblem{:nonsymmetric,:row}, |
| 89 | + ::StructuredColoringAlgorithm; |
| 90 | + kwargs..., |
| 91 | +) |
| 92 | + color = blockbanded_coloring(A, 1) |
| 93 | + bg = BipartiteGraph(A) |
| 94 | + return RowColoringResult(A, bg, color) |
| 95 | +end |
| 96 | + |
| 97 | +end |
0 commit comments