SparseMatrixColorings.jl/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl at b24c181985a101be4682e397b4c9d3cbea29f7f7 · JuliaDiff/SparseMatrixColorings.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
module SparseMatrixColoringsBlockBandedMatricesExt

using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
using BlockBandedMatrices:
    BandedBlockBandedMatrix,
    BlockBandedMatrix,
    blockbandrange,
    blockbandwidths,
    blocklengths,
    blocksize,
    subblockbandwidths
using SparseMatrixColorings:
    BipartiteGraph,
    ColoringProblem,
    ColumnColoringResult,
    StructuredColoringAlgorithm,
    RowColoringResult,
    column_colors,
    cycle_range,
    row_colors
import SparseMatrixColorings as SMC

#=
This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
https://github.com/JuliaArrays/ArrayInterface.jl
https://github.com/JuliaDiff/FiniteDiff.jl
=#

function subblockbandrange(A::BandedBlockBandedMatrix)
    u, l = subblockbandwidths(A)
    return (-l):u
end

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))
    first_in_block = blockfirsts(axes(A, dim))
    last_in_block = blocklasts(axes(A, dim))
    color = zeros(Int, size(A, dim))

    # give a macroscopic color to each block, so that 2 blocks with the same macro color are orthogonal
    # same idea as for BandedMatrices
    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)
        typemax(Int)
    end

    # for each macroscopic color, count how many microscopic colors will be needed
    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
        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::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
    ::ColoringProblem{:nonsymmetric,:column},
    ::StructuredColoringAlgorithm;
    kwargs...,
)
    color = blockbanded_coloring(A, 2)
    bg = BipartiteGraph(A)
    return ColumnColoringResult(A, bg, color)
end

function SMC.coloring(
    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
    ::ColoringProblem{:nonsymmetric,:row},
    ::StructuredColoringAlgorithm;
    kwargs...,
)
    color = blockbanded_coloring(A, 1)
    bg = BipartiteGraph(A)
    return RowColoringResult(A, bg, color)
end

end