diff --git a/Project.toml b/Project.toml index 9ef7826d..72a5fad2 100644 --- a/Project.toml +++ b/Project.toml @@ -12,6 +12,9 @@ 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" @@ -21,6 +24,8 @@ MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" cuSPARSE = "b26da814-b3bc-49ef-b0ee-c816305aa060" [extensions] +SparseMatrixColoringsBandedMatricesExt = "BandedMatrices" +SparseMatrixColoringsBlockBandedMatricesExt = ["BlockArrays", "BlockBandedMatrices"] SparseMatrixColoringsCUDAExt = ["CUDA", "cuSPARSE"] SparseMatrixColoringsCliqueTreesExt = "CliqueTrees" SparseMatrixColoringsColorsExt = "Colors" @@ -29,6 +34,9 @@ SparseMatrixColoringsJuMPExt = ["JuMP", "MathOptInterface"] [compat] ADTypes = "1.2.1" +BandedMatrices = "1.9.4" +BlockArrays = "1.6.3" +BlockBandedMatrices = "0.13.1" CUDA = "6.0.0" CliqueTrees = "1" Colors = "0.12.11, 0.13" diff --git a/docs/make.jl b/docs/make.jl index 80d2a23d..bbf3f859 100644 --- a/docs/make.jl +++ b/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) diff --git a/docs/src/api.md b/docs/src/api.md index f90c888f..640198d3 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -25,6 +25,7 @@ ColoringProblem GreedyColoringAlgorithm ConstantColoringAlgorithm OptimalColoringAlgorithm +StructuredColoringAlgorithm ``` ## Result analysis diff --git a/docs/src/dev.md b/docs/src/dev.md index 0f7b86fa..cd294a17 100644 --- a/docs/src/dev.md +++ b/docs/src/dev.md @@ -80,3 +80,9 @@ SparseMatrixColorings.what_fig_61 SparseMatrixColorings.efficient_fig_1 SparseMatrixColorings.efficient_fig_4 ``` + +## Misc + +```@docs +SparseMatrixColorings.cycle_range +``` diff --git a/ext/SparseMatrixColoringsBandedMatricesExt.jl b/ext/SparseMatrixColoringsBandedMatricesExt.jl new file mode 100644 index 00000000..05d9677c --- /dev/null +++ b/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 diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl new file mode 100644 index 00000000..aaecb0ad --- /dev/null +++ b/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 diff --git a/ext/SparseMatrixColoringsCUDAExt.jl b/ext/SparseMatrixColoringsCUDAExt.jl index ed33eece..59b379a9 100644 --- a/ext/SparseMatrixColoringsCUDAExt.jl +++ b/ext/SparseMatrixColoringsCUDAExt.jl @@ -1,10 +1,35 @@ module SparseMatrixColoringsCUDAExt - +using LinearAlgebra import SparseMatrixColorings as SMC using SparseArrays: SparseMatrixCSC, rowvals, nnz, nzrange -using CUDA: CuVector, CuMatrix +using CUDA: CuArray, CuVector, CuMatrix using cuSPARSE: AbstractCuSparseMatrix, CuSparseMatrixCSC, CuSparseMatrixCSR +## Basic support for GPU sparsity pattern stuff + +function SMC.SparsityPatternCSC(A::CuSparseMatrixCSC) + SMC.SparsityPatternCSC(first(A.dims), last(A.dims), A.colPtr, A.rowVal) +end + +for R in (:Diagonal, :Bidiagonal, :Tridiagonal) + @eval function SMC.BipartiteGraph( + A::$R{T,<:CuArray}; symmetric_pattern::Bool=false + ) where {T} + return SMC.BipartiteGraph(CuSparseMatrixCSC(A); symmetric_pattern) + end +end + +function SMC.BipartiteGraph(A::CuSparseMatrixCSC; symmetric_pattern::Bool=false) + S2 = SMC.SparsityPatternCSC(A) + if symmetric_pattern + checksquare(A) # proxy for checking full symmetry + S1 = S2 + else + S1 = transpose(S2) # rows to columns + end + return SMC.BipartiteGraph(S1, S2) +end + ## CSC Result function SMC.ColumnColoringResult( diff --git a/ext/SparseMatrixColoringsColorsExt.jl b/ext/SparseMatrixColoringsColorsExt.jl index c934b5e9..7f844e89 100644 --- a/ext/SparseMatrixColoringsColorsExt.jl +++ b/ext/SparseMatrixColoringsColorsExt.jl @@ -271,8 +271,9 @@ function show_colors!( A_ccolor_indices = mod1.(column_colors(res), length(colorscheme)) A_rcolor_indices = mod1.(row_shift .+ row_colors(res), length(colorscheme)) B_ccolor_indices = mod1.(1:maximum(column_colors(res)), length(colorscheme)) - B_rcolor_indices = - mod1.((row_shift + 1):(row_shift + maximum(row_colors(res))), length(colorscheme)) + B_rcolor_indices = mod1.( + (row_shift + 1):(row_shift + maximum(row_colors(res))), length(colorscheme) + ) A_ccolors = colorscheme[A_ccolor_indices] A_rcolors = colorscheme[A_rcolor_indices] B_ccolors = colorscheme[B_ccolor_indices] diff --git a/ext/SparseMatrixColoringsJuMPExt.jl b/ext/SparseMatrixColoringsJuMPExt.jl index 77b74dbd..7e5c7ff3 100644 --- a/ext/SparseMatrixColoringsJuMPExt.jl +++ b/ext/SparseMatrixColoringsJuMPExt.jl @@ -29,7 +29,7 @@ function optimal_distance2_coloring( model = Model(optimizer) silent && set_silent(model) # one variable per vertex to color, removing some renumbering symmetries - @variable(model, 1 <= color[i=1:n] <= i, Int) + @variable(model, 1 <= color[i = 1:n] <= i, Int) # one variable to count the number of distinct colors @variable(model, ncolors, Int) @constraint(model, [ncolors; color] in MOI.CountDistinct(n + 1)) diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl index 481c8584..74486d80 100644 --- a/src/SparseMatrixColorings.jl +++ b/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, @@ -54,6 +56,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") @@ -65,7 +68,7 @@ export NaturalOrder, RandomOrder, LargestFirst export DynamicDegreeBasedOrder, SmallestLast, IncidenceDegree, DynamicLargestFirst export PerfectEliminationOrder export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult -export ConstantColoringAlgorithm +export ConstantColoringAlgorithm, StructuredColoringAlgorithm export OptimalColoringAlgorithm export coloring, fast_coloring export column_colors, row_colors, ncolors diff --git a/src/graph.jl b/src/graph.jl index 94954602..adc33561 100644 --- a/src/graph.jl +++ b/src/graph.jl @@ -17,8 +17,8 @@ Copied from `SparseMatrixCSC`: struct SparsityPatternCSC{Ti<:Integer} <: AbstractMatrix{Bool} m::Int n::Int - colptr::Vector{Ti} - rowval::Vector{Ti} + colptr::AbstractVector{Ti} + rowval::AbstractVector{Ti} end SparsityPatternCSC(A::SparseMatrixCSC) = SparsityPatternCSC(A.m, A.n, A.colptr, A.rowval) diff --git a/src/structured.jl b/src/structured.jl new file mode 100644 index 00000000..812256bf --- /dev/null +++ b/src/structured.jl @@ -0,0 +1,181 @@ +""" + StructuredColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm + +Coloring algorithm which leverages specific matrix structures to produce optimal or near-optimal solutions. + +The following matrix types are supported: + +- From the standard library `LinearAlgebra`: `Diagonal`, `Bidiagonal`, `Tridiagonal` +- From [BandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BandedMatrices.jl): [`BandedMatrix`](@extref BandedMatrices.BandedMatrix) + +!!! warning + Only `:nonsymmetric` structures with `:row` or `:column` partitions (aka unidirectional Jacobian colorings) are supported by this algorithm at the moment. + +!!! tip + To request support for a new type of structured matrix, open an issue on the SparseMatrixColorings.jl GitHub repository! +""" +struct StructuredColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm end + +#= +This code is partly taken from ArrayInterface.jl +https://github.com/JuliaArrays/ArrayInterface.jl +=# + +""" + cycle_range(k::Integer, n::Integer) + +Concatenate copies of `1:k` to fill a vector of length `n` (with one partial copy allowed at the end). +""" +function cycle_range(k::Integer, n::Integer) + color = Vector{Int}(undef, n) + for i in eachindex(color) + color[i] = 1 + (i - 1) % k + end + return color +end + +## Diagonal + +function coloring( + A::Diagonal, + ::ColoringProblem{:nonsymmetric,:column}, + ::StructuredColoringAlgorithm; + kwargs..., +) + color = fill(1, size(A, 2)) + bg = BipartiteGraph(A) + return ColumnColoringResult(A, bg, color) +end + +function coloring( + A::Diagonal, + ::ColoringProblem{:nonsymmetric,:row}, + ::StructuredColoringAlgorithm; + kwargs..., +) + color = fill(1, size(A, 1)) + bg = BipartiteGraph(A) + return RowColoringResult(A, bg, color) +end + +function decompress!(A::Diagonal, B::AbstractMatrix, result::ColumnColoringResult) + color = column_colors(result) + for j in axes(A, 2) + A[j, j] = B[j, color[j]] + end + return A +end + +function decompress!(A::Diagonal, B::AbstractMatrix, result::RowColoringResult) + color = row_colors(result) + for i in axes(A, 1) + A[i, i] = B[color[i], i] + end + return A +end + +## Bidiagonal + +function coloring( + A::Bidiagonal, + ::ColoringProblem{:nonsymmetric,:column}, + ::StructuredColoringAlgorithm; + kwargs..., +) + color = cycle_range(2, size(A, 2)) + bg = BipartiteGraph(A) + return ColumnColoringResult(A, bg, color) +end + +function coloring( + A::Bidiagonal, + ::ColoringProblem{:nonsymmetric,:row}, + ::StructuredColoringAlgorithm; + kwargs..., +) + color = cycle_range(2, size(A, 1)) + bg = BipartiteGraph(A) + return RowColoringResult(A, bg, color) +end + +function decompress!(A::Bidiagonal, B::AbstractMatrix, result::ColumnColoringResult) + color = column_colors(result) + for j in axes(A, 2) + c = color[j] + A[j, j] = B[j, c] + if A.uplo == 'U' && j > 1 # above + A[j - 1, j] = B[j - 1, c] + elseif A.uplo == 'L' && j < size(A, 2) # below + A[j + 1, j] = B[j + 1, c] + end + end + return A +end + +function decompress!(A::Bidiagonal, B::AbstractMatrix, result::RowColoringResult) + color = row_colors(result) + for i in axes(A, 1) + c = color[i] + A[i, i] = B[c, i] + if A.uplo == 'U' && i < size(A, 1) # right + A[i, i + 1] = B[c, i + 1] + elseif A.uplo == 'L' && i > 1 # left + A[i, i - 1] = B[c, i - 1] + end + end + return A +end + +## Tridiagonal + +function coloring( + A::Tridiagonal, + ::ColoringProblem{:nonsymmetric,:column}, + ::StructuredColoringAlgorithm; + kwargs..., +) + color = cycle_range(3, size(A, 2)) + bg = BipartiteGraph(A) + return ColumnColoringResult(A, bg, color) +end + +function coloring( + A::Tridiagonal, + ::ColoringProblem{:nonsymmetric,:row}, + ::StructuredColoringAlgorithm; + kwargs..., +) + color = cycle_range(3, size(A, 1)) + bg = BipartiteGraph(A) + return RowColoringResult(A, bg, color) +end + +function decompress!(A::Tridiagonal, B::AbstractMatrix, result::ColumnColoringResult) + color = column_colors(result) + for j in axes(A, 2) + c = color[j] + A[j, j] = B[j, c] + if j > 1 # above + A[j - 1, j] = B[j - 1, c] + end + if j < size(A, 2) # below + A[j + 1, j] = B[j + 1, c] + end + end + return A +end + +function decompress!(A::Tridiagonal, B::AbstractMatrix, result::RowColoringResult) + color = row_colors(result) + for i in axes(A, 1) + c = color[i] + A[i, i] = B[c, i] + if i < size(A, 1) # right + A[i, i + 1] = B[c, i + 1] + end + if i > 1 # left + A[i, i - 1] = B[c, i - 1] + end + end + return A +end diff --git a/test/structured.jl b/test/structured.jl index 0d2a7564..857ebbfc 100644 --- a/test/structured.jl +++ b/test/structured.jl @@ -7,54 +7,66 @@ using SparseMatrixColorings using Test @testset "Diagonal" begin - for n in (1, 2, 10, 100) - A = Diagonal(rand(n)) - test_structured_coloring_decompression(A) + @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()] + for n in (1, 2, 10, 100) + A = Diagonal(rand(n)) + test_structured_coloring_decompression(A, algo) + end 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) + @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()] + 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, algo) + test_structured_coloring_decompression(A2, algo) + end 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) + @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()] + for n in (2, 10, 100) + A = Tridiagonal(rand(n - 1), rand(n), rand(n - 1)) + test_structured_coloring_decompression(A, algo) + end 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) + @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()] + @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, algo) + end 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) + @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()] + 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, algo) + end 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) + @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()] + 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, algo) + end end end; diff --git a/test/utils.jl b/test/utils.jl index 77c88610..15aa3f9f 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -274,10 +274,11 @@ function test_bicoloring_decompression( end end -function test_structured_coloring_decompression(A::AbstractMatrix) +function test_structured_coloring_decompression( + A::AbstractMatrix, algo=StructuredColoringAlgorithm() +) column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column) row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row) - algo = GreedyColoringAlgorithm() # Column result = coloring(A, column_problem, algo) @@ -287,9 +288,8 @@ function test_structured_coloring_decompression(A::AbstractMatrix) @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 + if algo isa StructuredColoringAlgorithm + @test color == ArrayInterface.matrix_colors(A) end # Row