Optimized implementation for structured matrices

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseMatrixColorings"
 uuid = "0a514795-09f3-496d-8182-132a7b665d35"
 authors = ["Guillaume Dalle", "Alexis Montoison"]
-version = "0.4.21"
+version = "0.4.22"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -12,17 +12,25 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 
 [extensions]
+SparseMatrixColoringsBandedMatricesExt = "BandedMatrices"
+SparseMatrixColoringsBlockBandedMatricesExt = ["BlockArrays", "BlockBandedMatrices"]
 SparseMatrixColoringsCUDAExt = "CUDA"
 SparseMatrixColoringsCliqueTreesExt = "CliqueTrees"
 SparseMatrixColoringsColorsExt = "Colors"
 
 [compat]
 ADTypes = "1.2.1"
+BandedMatrices = "1.9.4"
+BlockArrays = "1.6.3"
+BlockBandedMatrices = "0.13.1"
 CUDA = "5.8.2"
 CliqueTrees = "1"
 Colors = "0.12.11, 0.13"

docs/make.jl

@@ -2,7 +2,10 @@ using Documenter
 using DocumenterInterLinks
 using SparseMatrixColorings
 
-links = InterLinks("ADTypes" => "https://sciml.github.io/ADTypes.jl/stable/")
+links = InterLinks(
+    "ADTypes" => "https://sciml.github.io/ADTypes.jl/stable/",
+    "BandedMatrices" => "https://julialinearalgebra.github.io/BandedMatrices.jl/stable/",
+)
 
 cp(joinpath(@__DIR__, "..", "README.md"), joinpath(@__DIR__, "src", "index.md"); force=true)
 

docs/src/api.md

@@ -17,8 +17,14 @@ SparseMatrixColorings
 coloring
 fast_coloring
 ColoringProblem
+```
+
+## Coloring algorithms
+
+```@docs
 GreedyColoringAlgorithm
 ConstantColoringAlgorithm
+StructuredColoringAlgorithm
 ```
 
 ## Result analysis

docs/src/dev.md

@@ -76,3 +76,9 @@ SparseMatrixColorings.what_fig_61
 SparseMatrixColorings.efficient_fig_1
 SparseMatrixColorings.efficient_fig_4
 ```
+
+## Misc
+
+```@docs
+SparseMatrixColorings.cycle_range
+```

ext/SparseMatrixColoringsBandedMatricesExt.jl

@@ -0,0 +1,45 @@
+module SparseMatrixColoringsBandedMatricesExt
+
+using BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
+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 SMC.coloring(
+    A::BandedMatrix,
+    ::ColoringProblem{:nonsymmetric,:column},
+    ::StructuredColoringAlgorithm;
+    kwargs...,
+)
+    width = length(bandrange(A))
+    color = cycle_range(width, size(A, 2))
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
+end
+
+function SMC.coloring(
+    A::BandedMatrix,
+    ::ColoringProblem{:nonsymmetric,:row},
+    ::StructuredColoringAlgorithm;
+    kwargs...,
+)
+    width = length(bandrange(A))
+    color = cycle_range(width, size(A, 1))
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
+end
+
+end

ext/SparseMatrixColoringsBlockBandedMatricesExt.jl

@@ -0,0 +1,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

src/SparseMatrixColorings.jl

@@ -14,11 +14,13 @@ using Base.Iterators: Iterators
 using DocStringExtensions: README, EXPORTS, SIGNATURES, TYPEDEF, TYPEDFIELDS
 using LinearAlgebra:
     Adjoint,
+    Bidiagonal,
     Diagonal,
     Hermitian,
     LowerTriangular,
     Symmetric,
     Transpose,
+    Tridiagonal,
     UpperTriangular,
     adjoint,
     checksquare,
@@ -53,6 +55,7 @@ include("interface.jl")
 include("constant.jl")
 include("adtypes.jl")
 include("decompression.jl")
+include("structured.jl")
 include("check.jl")
 include("examples.jl")
 include("show_colors.jl")
@@ -63,7 +66,7 @@ export NaturalOrder, RandomOrder, LargestFirst
 export DynamicDegreeBasedOrder, SmallestLast, IncidenceDegree, DynamicLargestFirst
 export PerfectEliminationOrder
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
-export ConstantColoringAlgorithm
+export ConstantColoringAlgorithm, StructuredColoringAlgorithm
 export coloring, fast_coloring
 export column_colors, row_colors, ncolors
 export column_groups, row_groups