Make any coloring algorithm compatible with SMC

.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

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

src/coloring.jl

@@ -1,12 +1,19 @@
+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.
 
 A _distance-2 coloring_ is such that two vertices have different colors if they are at distance at most 2.
 
 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

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