Skip to content

Commit b53dbf9

Browse files
authored
Make any coloring algorithm compatible with SMC (#263)
* Make any coloring algorithm compatible with SMC * Fix StackOverflow * Run tests on 1.11 * Start working on optimal coloring * Fix tests * Improve feasibility check * Format * Fix * Fix * Typo * Use HiGHS * Remove OptimalColoringAlgorithm * Remove test deps * Don't handle forced colors in acyclic * Run GPU CI on 1.11 too * Actually use forced colors * Format * Deactivate JuliaFormatter (can't figure out why it fails)
1 parent e77226a commit b53dbf9

File tree

9 files changed

+226
-119
lines changed

9 files changed

+226
-119
lines changed

.github/workflows/Test-GPU.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ jobs:
2424
JULIA_SMC_TEST_GROUP: "GPU"
2525
strategy:
2626
matrix:
27-
julia-version: ['1.10', '1']
27+
julia-version: ['1.10', '1.11']
2828
steps:
2929
- uses: actions/checkout@v5
3030
- uses: julia-actions/setup-julia@v2

src/adtypes.jl

Lines changed: 29 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,33 @@
1-
function coloring(
2-
A::AbstractMatrix,
3-
::ColoringProblem{:nonsymmetric,:column},
4-
algo::ADTypes.NoColoringAlgorithm;
5-
kwargs...,
6-
)
7-
bg = BipartiteGraph(A)
8-
color = convert(Vector{eltype(bg)}, ADTypes.column_coloring(A, algo))
9-
return ColumnColoringResult(A, bg, color)
10-
end
1+
## From ADTypes to SMC
112

123
function coloring(
134
A::AbstractMatrix,
14-
::ColoringProblem{:nonsymmetric,:row},
15-
algo::ADTypes.NoColoringAlgorithm;
16-
kwargs...,
17-
)
18-
bg = BipartiteGraph(A)
19-
color = convert(Vector{eltype(bg)}, ADTypes.row_coloring(A, algo))
20-
return RowColoringResult(A, bg, color)
5+
problem::ColoringProblem{structure,partition},
6+
algo::ADTypes.AbstractColoringAlgorithm;
7+
decompression_eltype::Type{R}=Float64,
8+
symmetric_pattern::Bool=false,
9+
) where {structure,partition,R}
10+
symmetric_pattern = symmetric_pattern || A isa Union{Symmetric,Hermitian}
11+
if structure == :nonsymmetric
12+
if partition == :column
13+
forced_colors = ADTypes.column_coloring(A, algo)
14+
elseif partition == :row
15+
forced_colors = ADTypes.row_coloring(A, algo)
16+
else
17+
# TODO: improve once https://github.com/SciML/ADTypes.jl/issues/69 is done
18+
A_and_Aᵀ, _ = bidirectional_pattern(A; symmetric_pattern)
19+
forced_colors = ADTypes.symmetric_coloring(A_and_Aᵀ, algo)
20+
end
21+
else
22+
forced_colors = ADTypes.symmetric_coloring(A, algo)
23+
end
24+
return _coloring(
25+
WithResult(),
26+
A,
27+
problem,
28+
GreedyColoringAlgorithm(),
29+
R,
30+
symmetric_pattern;
31+
forced_colors,
32+
)
2133
end

src/coloring.jl

Lines changed: 56 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,19 @@
1+
struct InvalidColoringError <: Exception end
2+
13
"""
2-
partial_distance2_coloring(bg::BipartiteGraph, ::Val{side}, vertices_in_order::AbstractVector)
4+
partial_distance2_coloring(
5+
bg::BipartiteGraph, ::Val{side}, vertices_in_order::AbstractVector;
6+
forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing
7+
)
38
49
Compute a distance-2 coloring of the given `side` (`1` or `2`) in the bipartite graph `bg` and return a vector of integer colors.
510
611
A _distance-2 coloring_ is such that two vertices have different colors if they are at distance at most 2.
712
813
The vertices are colored in a greedy fashion, following the order supplied.
914
15+
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.
16+
1017
# See also
1118
1219
- [`BipartiteGraph`](@ref)
@@ -17,11 +24,16 @@ The vertices are colored in a greedy fashion, following the order supplied.
1724
> [_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
1825
"""
1926
function partial_distance2_coloring(
20-
bg::BipartiteGraph{T}, ::Val{side}, vertices_in_order::AbstractVector{<:Integer}
27+
bg::BipartiteGraph{T},
28+
::Val{side},
29+
vertices_in_order::AbstractVector{<:Integer};
30+
forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
2131
) where {T,side}
2232
color = Vector{T}(undef, nb_vertices(bg, Val(side)))
2333
forbidden_colors = Vector{T}(undef, nb_vertices(bg, Val(side)))
24-
partial_distance2_coloring!(color, forbidden_colors, bg, Val(side), vertices_in_order)
34+
partial_distance2_coloring!(
35+
color, forbidden_colors, bg, Val(side), vertices_in_order; forced_colors
36+
)
2537
return color
2638
end
2739

@@ -30,7 +42,8 @@ function partial_distance2_coloring!(
3042
forbidden_colors::AbstractVector{<:Integer},
3143
bg::BipartiteGraph,
3244
::Val{side},
33-
vertices_in_order::AbstractVector{<:Integer},
45+
vertices_in_order::AbstractVector{<:Integer};
46+
forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
3447
) where {side}
3548
color .= 0
3649
forbidden_colors .= 0
@@ -44,17 +57,32 @@ function partial_distance2_coloring!(
4457
end
4558
end
4659
end
47-
for i in eachindex(forbidden_colors)
48-
if forbidden_colors[i] != v
49-
color[v] = i
50-
break
60+
if isnothing(forced_colors)
61+
for i in eachindex(forbidden_colors)
62+
if forbidden_colors[i] != v
63+
color[v] = i
64+
break
65+
end
66+
end
67+
else
68+
f = forced_colors[v]
69+
if (
70+
(f == 0 && length(neighbors(bg, Val(side), v)) > 0) ||
71+
(f > 0 && forbidden_colors[f] == v)
72+
)
73+
throw(InvalidColoringError())
74+
else
75+
color[v] = f
5176
end
5277
end
5378
end
5479
end
5580

5681
"""
57-
star_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool)
82+
star_coloring(
83+
g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool;
84+
forced_colors::Union{AbstractVector,Nothing}=nothing
85+
)
5886
5987
Compute a star coloring of all vertices in the adjacency graph `g` and return a tuple `(color, star_set)`, where
6088
@@ -67,6 +95,8 @@ The vertices are colored in a greedy fashion, following the order supplied.
6795
6896
If `postprocessing=true`, some colors might be replaced with `0` (the "neutral" color) as long as they are not needed during decompression.
6997
98+
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.
99+
70100
# See also
71101
72102
- [`AdjacencyGraph`](@ref)
@@ -77,7 +107,10 @@ If `postprocessing=true`, some colors might be replaced with `0` (the "neutral"
77107
> [_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
78108
"""
79109
function star_coloring(
80-
g::AdjacencyGraph{T}, vertices_in_order::AbstractVector{<:Integer}, postprocessing::Bool
110+
g::AdjacencyGraph{T},
111+
vertices_in_order::AbstractVector{<:Integer},
112+
postprocessing::Bool;
113+
forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
81114
) where {T<:Integer}
82115
# Initialize data structures
83116
nv = nb_vertices(g)
@@ -115,10 +148,18 @@ function star_coloring(
115148
end
116149
end
117150
end
118-
for i in eachindex(forbidden_colors)
119-
if forbidden_colors[i] != v
120-
color[v] = i
121-
break
151+
if isnothing(forced_colors)
152+
for i in eachindex(forbidden_colors)
153+
if forbidden_colors[i] != v
154+
color[v] = i
155+
break
156+
end
157+
end
158+
else
159+
if forbidden_colors[forced_colors[v]] == v # TODO: handle forced_colors[v] == 0
160+
throw(InvalidColoringError())
161+
else
162+
color[v] = forced_colors[v]
122163
end
123164
end
124165
_update_stars!(star, hub, g, v, color, first_neighbor)
@@ -271,6 +312,7 @@ function acyclic_coloring(
271312
end
272313
end
273314
end
315+
# TODO: handle forced colors
274316
for i in eachindex(forbidden_colors)
275317
if forbidden_colors[i] != v
276318
color[v] = i

src/constant.jl

Lines changed: 46 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -4,20 +4,24 @@
44
Coloring algorithm which always returns the same precomputed vector of colors.
55
Useful when the optimal coloring of a matrix can be determined a priori due to its specific structure (e.g. banded).
66
7-
It is passed as an argument to the main function [`coloring`](@ref), but will only work if the associated `problem` has `:nonsymmetric` structure.
8-
Indeed, for symmetric coloring problems, we need more than just the vector of colors to allow fast decompression.
7+
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.
98
109
# Constructors
1110
1211
ConstantColoringAlgorithm{partition}(matrix_template, color)
13-
ConstantColoringAlgorithm(matrix_template, color; partition=:column)
12+
ConstantColoringAlgorithm{partition,structure}(matrix_template, color)
13+
ConstantColoringAlgorithm(
14+
matrix_template, color;
15+
structure=:nonsymmetric, partition=:column
16+
)
1417
1518
- `partition::Symbol`: either `:row` or `:column`.
19+
- `structure::Symbol`: either `:nonsymmetric` or `:symmetric`.
1620
- `matrix_template::AbstractMatrix`: matrix for which the vector of colors was precomputed (the algorithm will only accept matrices of the exact same size).
1721
- `color::Vector{<:Integer}`: vector of integer colors, one for each row or column (depending on `partition`).
1822
1923
!!! warning
20-
The second constructor (based on keyword arguments) is type-unstable.
24+
The constructor based on keyword arguments is type-unstable if these arguments are not compile-time constants.
2125
2226
We do not necessarily verify consistency between the matrix template and the vector of colors, this is the responsibility of the user.
2327
@@ -63,71 +67,68 @@ julia> column_colors(result)
6367
6468
- [`ADTypes.column_coloring`](@extref ADTypes.column_coloring)
6569
- [`ADTypes.row_coloring`](@extref ADTypes.row_coloring)
70+
- [`ADTypes.symmetric_coloring`](@extref ADTypes.symmetric_coloring)
6671
"""
67-
struct ConstantColoringAlgorithm{
68-
partition,
69-
M<:AbstractMatrix,
70-
T<:Integer,
71-
R<:AbstractColoringResult{:nonsymmetric,partition,:direct},
72-
} <: ADTypes.AbstractColoringAlgorithm
72+
struct ConstantColoringAlgorithm{partition,structure,M<:AbstractMatrix,T<:Integer} <:
73+
ADTypes.AbstractColoringAlgorithm
7374
matrix_template::M
7475
color::Vector{T}
75-
result::R
76-
end
7776

78-
function ConstantColoringAlgorithm{:column}(
79-
matrix_template::AbstractMatrix, color::Vector{<:Integer}
80-
)
81-
bg = BipartiteGraph(matrix_template)
82-
result = ColumnColoringResult(matrix_template, bg, color)
83-
T, M, R = eltype(bg), typeof(matrix_template), typeof(result)
84-
return ConstantColoringAlgorithm{:column,M,T,R}(matrix_template, color, result)
77+
function ConstantColoringAlgorithm{partition,structure}(
78+
matrix_template::AbstractMatrix, color::Vector{<:Integer}
79+
) where {partition,structure}
80+
check_valid_problem(structure, partition)
81+
return new{partition,structure,typeof(matrix_template),eltype(color)}(
82+
matrix_template, color
83+
)
84+
end
8585
end
8686

87-
function ConstantColoringAlgorithm{:row}(
87+
function ConstantColoringAlgorithm{partition}(
8888
matrix_template::AbstractMatrix, color::Vector{<:Integer}
89-
)
90-
bg = BipartiteGraph(matrix_template)
91-
result = RowColoringResult(matrix_template, bg, color)
92-
T, M, R = eltype(bg), typeof(matrix_template), typeof(result)
93-
return ConstantColoringAlgorithm{:row,M,T,R}(matrix_template, color, result)
89+
) where {partition}
90+
return ConstantColoringAlgorithm{partition,:nonsymmetric}(matrix_template, color)
9491
end
9592

9693
function ConstantColoringAlgorithm(
97-
matrix_template::AbstractMatrix, color::Vector{<:Integer}; partition::Symbol=:column
94+
matrix_template::AbstractMatrix,
95+
color::Vector{<:Integer};
96+
structure::Symbol=:nonsymmetric,
97+
partition::Symbol=:column,
9898
)
99-
return ConstantColoringAlgorithm{partition}(matrix_template, color)
99+
return ConstantColoringAlgorithm{partition,structure}(matrix_template, color)
100100
end
101101

102-
function coloring(
103-
A::AbstractMatrix,
104-
::ColoringProblem{:nonsymmetric,partition},
105-
algo::ConstantColoringAlgorithm{partition};
106-
decompression_eltype::Type=Float64,
107-
symmetric_pattern::Bool=false,
108-
) where {partition}
109-
(; matrix_template, result) = algo
102+
function check_template(algo::ConstantColoringAlgorithm, A::AbstractMatrix)
103+
(; matrix_template) = algo
110104
if size(A) != size(matrix_template)
111105
throw(
112106
DimensionMismatch(
113107
"`ConstantColoringAlgorithm` expected matrix of size $(size(matrix_template)) but got matrix of size $(size(A))",
114108
),
115109
)
116-
else
117-
return result
118110
end
119111
end
120112

121113
function ADTypes.column_coloring(
122-
A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column}
114+
A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column,:nonsymmetric}
115+
)
116+
check_template(algo, A)
117+
return algo.color
118+
end
119+
120+
function ADTypes.row_coloring(
121+
A::AbstractMatrix, algo::ConstantColoringAlgorithm{:row,:nonsymmetric}
123122
)
124-
problem = ColoringProblem{:nonsymmetric,:column}()
125-
result = coloring(A, problem, algo)
126-
return column_colors(result)
123+
check_template(algo, A)
124+
return algo.color
127125
end
128126

129-
function ADTypes.row_coloring(A::AbstractMatrix, algo::ConstantColoringAlgorithm)
130-
problem = ColoringProblem{:nonsymmetric,:row}()
131-
result = coloring(A, problem, algo)
132-
return row_colors(result)
127+
function ADTypes.symmetric_coloring(
128+
A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column,:symmetric}
129+
)
130+
check_template(algo, A)
131+
return algo.color
133132
end
133+
134+
# TODO: handle bidirectional once https://github.com/SciML/ADTypes.jl/issues/69 is done

0 commit comments

Comments
 (0)