Make any coloring algorithm compatible with SMC

Author: gdalle

src/adtypes.jl

@@ -1,21 +1,32 @@
-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
+            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,17 @@
 """
-    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 +22,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 +40,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 +55,25 @@ 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
+            @assert forbidden_colors[forced_colors[v]] != v
+            color[v] = forced_colors[v]
         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 +86,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 +98,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,11 +139,16 @@ 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
+            @assert forbidden_colors[forced_colors[v]] != v
+            color[v] = forced_colors[v]
         end
         _update_stars!(star, hub, g, v, color, first_neighbor)
     end
@@ -209,7 +238,10 @@ struct StarSet{T}
 end
 
 """
-    acyclic_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool)
+    acyclic_coloring(
+        g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool;
+        forced_colors::Union{AbstractVector,Nothing}=nothing
+    )
 
 Compute an acyclic coloring of all vertices in the adjacency graph `g` and return a tuple `(color, tree_set)`, where
 
@@ -222,6 +254,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 acyclic coloring procedure to verify correctness and build auxiliary data structures, useful during decompression.
+
 # See also
 
 - [`AdjacencyGraph`](@ref)
@@ -232,7 +266,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 3.1
 """
 function acyclic_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)
@@ -271,11 +308,16 @@ function acyclic_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
+            @assert forbidden_colors[forced_colors[v]] != v
+            color[v] = forced_colors[v]
         end
         for (w, index_vw) in neighbors_with_edge_indices(g, v)  # grow two-colored stars around the vertex v
             !has_diagonal(g) || (v == w && continue)

src/interface.jl

@@ -225,12 +225,13 @@ 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)
     vertices_in_order = vertices(bg, Val(2), algo.order)
-    color = partial_distance2_coloring(bg, Val(2), vertices_in_order)
+    color = partial_distance2_coloring(bg, Val(2), vertices_in_order; forced_colors)
     if speed_setting isa WithResult
         return ColumnColoringResult(A, bg, color)
     else
@@ -244,12 +245,13 @@ 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)
     vertices_in_order = vertices(bg, Val(1), algo.order)
-    color = partial_distance2_coloring(bg, Val(1), vertices_in_order)
+    color = partial_distance2_coloring(bg, Val(1), vertices_in_order; forced_colors)
     if speed_setting isa WithResult
         return RowColoringResult(A, bg, color)
     else
@@ -263,11 +265,14 @@ 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)
     vertices_in_order = vertices(ag, algo.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
+    )
     if speed_setting isa WithResult
         return StarSetColoringResult(A, ag, color, star_set)
     else
@@ -281,11 +286,14 @@ function _coloring(
     ::ColoringProblem{:symmetric,:column},
     algo::GreedyColoringAlgorithm{:substitution},
     decompression_eltype::Type{R},
-    symmetric_pattern::Bool,
+    symmetric_pattern::Bool;
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {R}
     ag = AdjacencyGraph(A; has_diagonal=true)
     vertices_in_order = vertices(ag, algo.order)
-    color, tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
+    color, tree_set = acyclic_coloring(
+        ag, vertices_in_order, algo.postprocessing; forced_colors
+    )
     if speed_setting isa WithResult
         return TreeSetColoringResult(A, ag, color, tree_set, R)
     else
@@ -299,12 +307,15 @@ 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)
     vertices_in_order = vertices(ag, algo.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
+    )
     if speed_setting isa WithResult
         symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
         return BicoloringResult(A, ag, symmetric_result, R)
@@ -322,12 +333,15 @@ function _coloring(
     ::ColoringProblem{:nonsymmetric,:bidirectional},
     algo::GreedyColoringAlgorithm{:substitution},
     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)
     vertices_in_order = vertices(ag, algo.order)
-    color, tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
+    color, tree_set = acyclic_coloring(
+        ag, vertices_in_order, algo.postprocessing; forced_colors
+    )
     if speed_setting isa WithResult
         symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R)
         return BicoloringResult(A, ag, symmetric_result, R)

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