diff --git a/.github/workflows/Test-GPU.yml b/.github/workflows/Test-GPU.yml index 05628de0..f5a4fbc8 100644 --- a/.github/workflows/Test-GPU.yml +++ b/.github/workflows/Test-GPU.yml @@ -24,7 +24,7 @@ jobs: JULIA_SMC_TEST_GROUP: "GPU" strategy: matrix: - julia-version: ['1.10', '1'] + julia-version: ['1.10', '1.11'] steps: - uses: actions/checkout@v5 - uses: julia-actions/setup-julia@v2 diff --git a/src/adtypes.jl b/src/adtypes.jl index 7500b1e2..7c464433 100644 --- a/src/adtypes.jl +++ b/src/adtypes.jl @@ -1,21 +1,33 @@ -function coloring( - A::AbstractMatrix, - ::ColoringProblem{:nonsymmetric,:column}, - algo::ADTypes.NoColoringAlgorithm; - kwargs..., -) - bg = BipartiteGraph(A) - color = convert(Vector{eltype(bg)}, ADTypes.column_coloring(A, algo)) - return ColumnColoringResult(A, bg, color) -end +## From ADTypes to SMC function coloring( A::AbstractMatrix, - ::ColoringProblem{:nonsymmetric,:row}, - algo::ADTypes.NoColoringAlgorithm; - kwargs..., -) - bg = BipartiteGraph(A) - color = convert(Vector{eltype(bg)}, ADTypes.row_coloring(A, algo)) - return RowColoringResult(A, bg, color) + problem::ColoringProblem{structure,partition}, + algo::ADTypes.AbstractColoringAlgorithm; + decompression_eltype::Type{R}=Float64, + symmetric_pattern::Bool=false, +) where {structure,partition,R} + symmetric_pattern = symmetric_pattern || A isa Union{Symmetric,Hermitian} + if structure == :nonsymmetric + if partition == :column + forced_colors = ADTypes.column_coloring(A, algo) + elseif partition == :row + forced_colors = ADTypes.row_coloring(A, algo) + else + # TODO: improve once https://github.com/SciML/ADTypes.jl/issues/69 is done + A_and_Aᵀ, _ = bidirectional_pattern(A; symmetric_pattern) + forced_colors = ADTypes.symmetric_coloring(A_and_Aᵀ, algo) + end + else + forced_colors = ADTypes.symmetric_coloring(A, algo) + end + return _coloring( + WithResult(), + A, + problem, + GreedyColoringAlgorithm(), + R, + symmetric_pattern; + forced_colors, + ) end diff --git a/src/coloring.jl b/src/coloring.jl index 7eec55c0..7288b773 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -1,5 +1,10 @@ +struct InvalidColoringError <: Exception end + """ - partial_distance2_coloring(bg::BipartiteGraph, ::Val{side}, vertices_in_order::AbstractVector) + partial_distance2_coloring( + bg::BipartiteGraph, ::Val{side}, vertices_in_order::AbstractVector; + forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing + ) Compute a distance-2 coloring of the given `side` (`1` or `2`) in the bipartite graph `bg` and return a vector of integer colors. @@ -7,6 +12,8 @@ A _distance-2 coloring_ is such that two vertices have different colors if they The vertices are colored in a greedy fashion, following the order supplied. +The optional `forced_colors` keyword argument is used to enforce predefined vertex colors (e.g. coming from another optimization algorithm) but still run the distance-2 coloring procedure to verify correctness. + # See also - [`BipartiteGraph`](@ref) @@ -17,11 +24,16 @@ The vertices are colored in a greedy fashion, following the order supplied. > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005), Algorithm 3.2 """ function partial_distance2_coloring( - bg::BipartiteGraph{T}, ::Val{side}, vertices_in_order::AbstractVector{<:Integer} + bg::BipartiteGraph{T}, + ::Val{side}, + vertices_in_order::AbstractVector{<:Integer}; + forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing, ) where {T,side} color = Vector{T}(undef, nb_vertices(bg, Val(side))) forbidden_colors = Vector{T}(undef, nb_vertices(bg, Val(side))) - partial_distance2_coloring!(color, forbidden_colors, bg, Val(side), vertices_in_order) + partial_distance2_coloring!( + color, forbidden_colors, bg, Val(side), vertices_in_order; forced_colors + ) return color end @@ -30,7 +42,8 @@ function partial_distance2_coloring!( forbidden_colors::AbstractVector{<:Integer}, bg::BipartiteGraph, ::Val{side}, - vertices_in_order::AbstractVector{<:Integer}, + vertices_in_order::AbstractVector{<:Integer}; + forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing, ) where {side} color .= 0 forbidden_colors .= 0 @@ -44,17 +57,32 @@ function partial_distance2_coloring!( end end end - for i in eachindex(forbidden_colors) - if forbidden_colors[i] != v - color[v] = i - break + if isnothing(forced_colors) + for i in eachindex(forbidden_colors) + if forbidden_colors[i] != v + color[v] = i + break + end + end + else + f = forced_colors[v] + if ( + (f == 0 && length(neighbors(bg, Val(side), v)) > 0) || + (f > 0 && forbidden_colors[f] == v) + ) + throw(InvalidColoringError()) + else + color[v] = f end end end end """ - star_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool) + star_coloring( + g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool; + forced_colors::Union{AbstractVector,Nothing}=nothing + ) Compute a star coloring of all vertices in the adjacency graph `g` and return a tuple `(color, star_set)`, where @@ -67,6 +95,8 @@ The vertices are colored in a greedy fashion, following the order supplied. If `postprocessing=true`, some colors might be replaced with `0` (the "neutral" color) as long as they are not needed during decompression. +The optional `forced_colors` keyword argument is used to enforce predefined vertex colors (e.g. coming from another optimization algorithm) but still run the star coloring procedure to verify correctness and build auxiliary data structures, useful during decompression. + # See also - [`AdjacencyGraph`](@ref) @@ -77,7 +107,10 @@ If `postprocessing=true`, some colors might be replaced with `0` (the "neutral" > [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 4.1 """ function star_coloring( - g::AdjacencyGraph{T}, vertices_in_order::AbstractVector{<:Integer}, postprocessing::Bool + g::AdjacencyGraph{T}, + vertices_in_order::AbstractVector{<:Integer}, + postprocessing::Bool; + forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing, ) where {T<:Integer} # Initialize data structures nv = nb_vertices(g) @@ -115,10 +148,18 @@ function star_coloring( end end end - for i in eachindex(forbidden_colors) - if forbidden_colors[i] != v - color[v] = i - break + if isnothing(forced_colors) + for i in eachindex(forbidden_colors) + if forbidden_colors[i] != v + color[v] = i + break + end + end + else + if forbidden_colors[forced_colors[v]] == v # TODO: handle forced_colors[v] == 0 + throw(InvalidColoringError()) + else + color[v] = forced_colors[v] end end _update_stars!(star, hub, g, v, color, first_neighbor) @@ -271,6 +312,7 @@ function acyclic_coloring( end end end + # TODO: handle forced colors for i in eachindex(forbidden_colors) if forbidden_colors[i] != v color[v] = i diff --git a/src/constant.jl b/src/constant.jl index d79caab0..56bf4101 100644 --- a/src/constant.jl +++ b/src/constant.jl @@ -4,20 +4,24 @@ Coloring algorithm which always returns the same precomputed vector of colors. Useful when the optimal coloring of a matrix can be determined a priori due to its specific structure (e.g. banded). -It is passed as an argument to the main function [`coloring`](@ref), but will only work if the associated `problem` has `:nonsymmetric` structure. -Indeed, for symmetric coloring problems, we need more than just the vector of colors to allow fast decompression. +It is passed as an argument to the main function [`coloring`](@ref), but will only work if the associated `problem` has a `:column` or `:row` partition. # Constructors ConstantColoringAlgorithm{partition}(matrix_template, color) - ConstantColoringAlgorithm(matrix_template, color; partition=:column) + ConstantColoringAlgorithm{partition,structure}(matrix_template, color) + ConstantColoringAlgorithm( + matrix_template, color; + structure=:nonsymmetric, partition=:column + ) - `partition::Symbol`: either `:row` or `:column`. +- `structure::Symbol`: either `:nonsymmetric` or `:symmetric`. - `matrix_template::AbstractMatrix`: matrix for which the vector of colors was precomputed (the algorithm will only accept matrices of the exact same size). - `color::Vector{<:Integer}`: vector of integer colors, one for each row or column (depending on `partition`). !!! warning - The second constructor (based on keyword arguments) is type-unstable. + The constructor based on keyword arguments is type-unstable if these arguments are not compile-time constants. We do not necessarily verify consistency between the matrix template and the vector of colors, this is the responsibility of the user. @@ -63,71 +67,68 @@ julia> column_colors(result) - [`ADTypes.column_coloring`](@extref ADTypes.column_coloring) - [`ADTypes.row_coloring`](@extref ADTypes.row_coloring) +- [`ADTypes.symmetric_coloring`](@extref ADTypes.symmetric_coloring) """ -struct ConstantColoringAlgorithm{ - partition, - M<:AbstractMatrix, - T<:Integer, - R<:AbstractColoringResult{:nonsymmetric,partition,:direct}, -} <: ADTypes.AbstractColoringAlgorithm +struct ConstantColoringAlgorithm{partition,structure,M<:AbstractMatrix,T<:Integer} <: + ADTypes.AbstractColoringAlgorithm matrix_template::M color::Vector{T} - result::R -end -function ConstantColoringAlgorithm{:column}( - matrix_template::AbstractMatrix, color::Vector{<:Integer} -) - bg = BipartiteGraph(matrix_template) - result = ColumnColoringResult(matrix_template, bg, color) - T, M, R = eltype(bg), typeof(matrix_template), typeof(result) - return ConstantColoringAlgorithm{:column,M,T,R}(matrix_template, color, result) + function ConstantColoringAlgorithm{partition,structure}( + matrix_template::AbstractMatrix, color::Vector{<:Integer} + ) where {partition,structure} + check_valid_problem(structure, partition) + return new{partition,structure,typeof(matrix_template),eltype(color)}( + matrix_template, color + ) + end end -function ConstantColoringAlgorithm{:row}( +function ConstantColoringAlgorithm{partition}( matrix_template::AbstractMatrix, color::Vector{<:Integer} -) - bg = BipartiteGraph(matrix_template) - result = RowColoringResult(matrix_template, bg, color) - T, M, R = eltype(bg), typeof(matrix_template), typeof(result) - return ConstantColoringAlgorithm{:row,M,T,R}(matrix_template, color, result) +) where {partition} + return ConstantColoringAlgorithm{partition,:nonsymmetric}(matrix_template, color) end function ConstantColoringAlgorithm( - matrix_template::AbstractMatrix, color::Vector{<:Integer}; partition::Symbol=:column + matrix_template::AbstractMatrix, + color::Vector{<:Integer}; + structure::Symbol=:nonsymmetric, + partition::Symbol=:column, ) - return ConstantColoringAlgorithm{partition}(matrix_template, color) + return ConstantColoringAlgorithm{partition,structure}(matrix_template, color) end -function coloring( - A::AbstractMatrix, - ::ColoringProblem{:nonsymmetric,partition}, - algo::ConstantColoringAlgorithm{partition}; - decompression_eltype::Type=Float64, - symmetric_pattern::Bool=false, -) where {partition} - (; matrix_template, result) = algo +function check_template(algo::ConstantColoringAlgorithm, A::AbstractMatrix) + (; matrix_template) = algo if size(A) != size(matrix_template) throw( DimensionMismatch( "`ConstantColoringAlgorithm` expected matrix of size $(size(matrix_template)) but got matrix of size $(size(A))", ), ) - else - return result end end function ADTypes.column_coloring( - A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column} + A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column,:nonsymmetric} +) + check_template(algo, A) + return algo.color +end + +function ADTypes.row_coloring( + A::AbstractMatrix, algo::ConstantColoringAlgorithm{:row,:nonsymmetric} ) - problem = ColoringProblem{:nonsymmetric,:column}() - result = coloring(A, problem, algo) - return column_colors(result) + check_template(algo, A) + return algo.color end -function ADTypes.row_coloring(A::AbstractMatrix, algo::ConstantColoringAlgorithm) - problem = ColoringProblem{:nonsymmetric,:row}() - result = coloring(A, problem, algo) - return row_colors(result) +function ADTypes.symmetric_coloring( + A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column,:symmetric} +) + check_template(algo, A) + return algo.color end + +# TODO: handle bidirectional once https://github.com/SciML/ADTypes.jl/issues/69 is done diff --git a/src/interface.jl b/src/interface.jl index c1a56c29..b8d4ce54 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -230,13 +230,14 @@ function _coloring( ::ColoringProblem{:nonsymmetric,:column}, algo::GreedyColoringAlgorithm, decompression_eltype::Type, - symmetric_pattern::Bool, + symmetric_pattern::Bool; + forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing, ) symmetric_pattern = symmetric_pattern || A isa Union{Symmetric,Hermitian} bg = BipartiteGraph(A; symmetric_pattern) color_by_order = map(algo.orders) do order vertices_in_order = vertices(bg, Val(2), order) - return partial_distance2_coloring(bg, Val(2), vertices_in_order) + return partial_distance2_coloring(bg, Val(2), vertices_in_order; forced_colors) end color = argmin(maximum, color_by_order) if speed_setting isa WithResult @@ -252,13 +253,14 @@ function _coloring( ::ColoringProblem{:nonsymmetric,:row}, algo::GreedyColoringAlgorithm, decompression_eltype::Type, - symmetric_pattern::Bool, + symmetric_pattern::Bool; + forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing, ) symmetric_pattern = symmetric_pattern || A isa Union{Symmetric,Hermitian} bg = BipartiteGraph(A; symmetric_pattern) color_by_order = map(algo.orders) do order vertices_in_order = vertices(bg, Val(1), order) - return partial_distance2_coloring(bg, Val(1), vertices_in_order) + return partial_distance2_coloring(bg, Val(1), vertices_in_order; forced_colors) end color = argmin(maximum, color_by_order) if speed_setting isa WithResult @@ -274,12 +276,13 @@ function _coloring( ::ColoringProblem{:symmetric,:column}, algo::GreedyColoringAlgorithm{:direct}, decompression_eltype::Type, - symmetric_pattern::Bool, + symmetric_pattern::Bool; + forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing, ) ag = AdjacencyGraph(A; has_diagonal=true) color_and_star_set_by_order = map(algo.orders) do order vertices_in_order = vertices(ag, order) - return star_coloring(ag, vertices_in_order, algo.postprocessing) + return star_coloring(ag, vertices_in_order, algo.postprocessing; forced_colors) end color, star_set = argmin(maximum ∘ first, color_and_star_set_by_order) if speed_setting isa WithResult @@ -316,13 +319,16 @@ function _coloring( ::ColoringProblem{:nonsymmetric,:bidirectional}, algo::GreedyColoringAlgorithm{:direct}, decompression_eltype::Type{R}, - symmetric_pattern::Bool, + symmetric_pattern::Bool; + forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing, ) where {R} A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern) ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false) outputs_by_order = map(algo.orders) do order vertices_in_order = vertices(ag, order) - _color, _star_set = star_coloring(ag, vertices_in_order, algo.postprocessing) + _color, _star_set = star_coloring( + ag, vertices_in_order, algo.postprocessing; forced_colors + ) (_row_color, _column_color, _symmetric_to_row, _symmetric_to_column) = remap_colors( eltype(ag), _color, maximum(_color), size(A)... ) diff --git a/test/adtypes.jl b/test/adtypes.jl index bbb124a6..da0807d3 100644 --- a/test/adtypes.jl +++ b/test/adtypes.jl @@ -1,26 +1,49 @@ using ADTypes: ADTypes using SparseArrays +using LinearAlgebra using SparseMatrixColorings using Test -@testset "Column coloring" begin - problem = ColoringProblem(; structure=:nonsymmetric, partition=:column) - algo = ADTypes.NoColoringAlgorithm() - A = sprand(10, 20, 0.1) - result = coloring(A, problem, algo) - B = compress(A, result) - @test size(B) == size(A) - @test column_colors(result) == ADTypes.column_coloring(A, algo) - @test decompress(B, result) == A -end +@testset "NoColoringAlgorithm" begin + @testset "Column coloring" begin + problem = ColoringProblem(; structure=:nonsymmetric, partition=:column) + algo = ADTypes.NoColoringAlgorithm() + A = sprand(10, 20, 0.3) + result = coloring(A, problem, algo) + B = compress(A, result) + @test size(B) == size(A) + @test column_colors(result) == ADTypes.column_coloring(A, algo) + @test decompress(B, result) == A + end + + @testset "Row coloring" begin + problem = ColoringProblem(; structure=:nonsymmetric, partition=:row) + algo = ADTypes.NoColoringAlgorithm() + A = sprand(10, 20, 0.3) + result = coloring(A, problem, algo) + B = compress(A, result) + @test size(B) == size(A) + @test row_colors(result) == ADTypes.row_coloring(A, algo) + @test decompress(B, result) == A + end + + @testset "Symmetric coloring" begin + problem = ColoringProblem(; structure=:symmetric, partition=:column) + algo = ADTypes.NoColoringAlgorithm() + A = Symmetric(sprand(20, 20, 0.3)) + result = coloring(A, problem, algo) + B = compress(A, result) + @test size(B) == size(A) + @test column_colors(result) == ADTypes.column_coloring(A, algo) + @test decompress(B, result) == A + end -@testset "Row coloring" begin - problem = ColoringProblem(; structure=:nonsymmetric, partition=:row) - algo = ADTypes.NoColoringAlgorithm() - A = sprand(10, 20, 0.1) - result = coloring(A, problem, algo) - B = compress(A, result) - @test size(B) == size(A) - @test row_colors(result) == ADTypes.row_coloring(A, algo) - @test decompress(B, result) == A + @testset "Bicoloring" begin + problem = ColoringProblem(; structure=:nonsymmetric, partition=:bidirectional) + algo = ADTypes.NoColoringAlgorithm() + A = sprand(10, 20, 0.3) + result = coloring(A, problem, algo) + Br, Bc = compress(A, result) + @test decompress(Br, Bc, result) == A + end end diff --git a/test/constant.jl b/test/constant.jl index 5de881b8..97c7f238 100644 --- a/test/constant.jl +++ b/test/constant.jl @@ -1,16 +1,20 @@ using ADTypes: ADTypes using SparseMatrixColorings +using SparseMatrixColorings: InvalidColoringError using Test -matrix_template = ones(100, 200) +matrix_template = ones(Bool, 10, 20) +sym_matrix_template = ones(Bool, 10, 10) @testset "Column coloring" begin problem = ColoringProblem(; structure=:nonsymmetric, partition=:column) - color = rand(1:5, size(matrix_template, 2)) + color = collect(1:20) algo = ConstantColoringAlgorithm(matrix_template, color; partition=:column) - wrong_algo = ConstantColoringAlgorithm(matrix_template, color; partition=:row) + wrong_algo = ConstantColoringAlgorithm{:row}(matrix_template, color) + wrong_color = ConstantColoringAlgorithm{:column}(matrix_template, ones(Int, 20)) @test_throws DimensionMismatch coloring(transpose(matrix_template), problem, algo) @test_throws MethodError coloring(matrix_template, problem, wrong_algo) + @test_throws InvalidColoringError coloring(matrix_template, problem, wrong_color) result = coloring(matrix_template, problem, algo) @test column_colors(result) == color @test ADTypes.column_coloring(matrix_template, algo) == color @@ -19,9 +23,13 @@ end @testset "Row coloring" begin problem = ColoringProblem(; structure=:nonsymmetric, partition=:row) - color = rand(1:5, size(matrix_template, 1)) + color = collect(1:10) algo = ConstantColoringAlgorithm(matrix_template, color; partition=:row) + wrong_algo = ConstantColoringAlgorithm{:column}(matrix_template, color) + wrong_color = ConstantColoringAlgorithm{:row}(matrix_template, ones(Int, 10)) @test_throws DimensionMismatch coloring(transpose(matrix_template), problem, algo) + @test_throws MethodError coloring(matrix_template, problem, wrong_algo) + @test_throws InvalidColoringError coloring(matrix_template, problem, wrong_color) result = coloring(matrix_template, problem, algo) @test row_colors(result) == color @test ADTypes.row_coloring(matrix_template, algo) == color @@ -29,9 +37,22 @@ end end @testset "Symmetric coloring" begin - wrong_problem = ColoringProblem(; structure=:symmetric, partition=:column) - color = rand(1:5, size(matrix_template, 2)) - algo = ConstantColoringAlgorithm(matrix_template, color; partition=:column) - @test_throws MethodError coloring(matrix_template, wrong_problem, algo) - @test_throws MethodError ADTypes.symmetric_coloring(matrix_template, algo) + problem = ColoringProblem(; structure=:symmetric, partition=:column) + color = collect(1:10) + algo = ConstantColoringAlgorithm( + sym_matrix_template, color; partition=:column, structure=:symmetric + ) + wrong_algo = ConstantColoringAlgorithm{:column,:nonsymmetric}( + sym_matrix_template, color + ) + wrong_color = ConstantColoringAlgorithm{:column,:symmetric}( + sym_matrix_template, ones(Int, 20) + ) + @test_throws DimensionMismatch coloring(matrix_template, problem, algo) + @test_throws MethodError coloring(sym_matrix_template, problem, wrong_algo) + @test_throws InvalidColoringError coloring(sym_matrix_template, problem, wrong_color) + result = coloring(sym_matrix_template, problem, algo) + @test column_colors(result) == color + @test ADTypes.symmetric_coloring(sym_matrix_template, algo) == color + @test_throws MethodError ADTypes.column_coloring(sym_matrix_template, algo) end diff --git a/test/result.jl b/test/result.jl index 0bea3353..49344aab 100644 --- a/test/result.jl +++ b/test/result.jl @@ -1,3 +1,4 @@ +using SparseMatrixColorings using SparseMatrixColorings: group_by_color, UnsupportedDecompressionError using Test @@ -20,7 +21,7 @@ using Test end @testset "Empty compression" begin - A = rand(10, 10) + A = zeros(Bool, 10, 10) color = zeros(Int, 10) problem = ColoringProblem{:nonsymmetric,:column}() algo = ConstantColoringAlgorithm(A, color; partition=:column) diff --git a/test/runtests.jl b/test/runtests.jl index 81c4ecc9..cfd51dce 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -24,11 +24,12 @@ include("utils.jl") @testset "JET" begin JET.test_package(SparseMatrixColorings; target_defined_modules=true) end - @testset "JuliaFormatter" begin - @test JuliaFormatter.format( - SparseMatrixColorings; verbose=false, overwrite=false - ) - end + # @testset "JuliaFormatter" begin + # TODO: switch to Runic (temporarily deactivated) + # @test JuliaFormatter.format( + # SparseMatrixColorings; verbose=false, overwrite=false + # ) + # end @testset "Doctests" begin Documenter.doctest(SparseMatrixColorings) end