Replace Coloring by SparseMatrixColoring

Project.toml

@@ -6,14 +6,14 @@ authors = ["Sophie Lequeu <slequeu@hotmail.com>", "Benoît Legat <benoit.legat@g
 [deps]
 Calculus = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [weakdeps]
@@ -27,7 +27,6 @@ ArrayDiffONNXExt = "ONNX"
 [compat]
 Calculus = "0.5.2"
 ChainRulesCore = "1"
-DataStructures = "0.18, 0.19"
 ForwardDiff = "1"
 JuMP = "1.29.4"
 LinearAlgebra = "1.10"
@@ -37,5 +36,6 @@ NaNMath = "1"
 ONNX = "0.2, 0.3"
 OrderedCollections = "1.8.1, 2"
 SparseArrays = "1.10"
+SparseMatrixColorings = "0.4.26"
 SpecialFunctions = "2.6.1"
 julia = "1.10"

src/ArrayDiff.jl

@@ -12,10 +12,11 @@ import LinearAlgebra
 import MathOptInterface as MOI
 const Nonlinear = MOI.Nonlinear
 import SparseArrays
+import SparseMatrixColorings as SMC
 import OrderedCollections
 
 """
-    Mode{S}() <: MOI.Nonlinear.AbstractAutomaticDifferentiation
+    Mode{S}([coloring_algorithm]) <: MOI.Nonlinear.AbstractAutomaticDifferentiation
 
 Fork of `MOI.Nonlinear.SparseReverseMode` to add array support.
 
@@ -24,9 +25,26 @@ partials, and reverse storage of each subexpression). It must satisfy
 `S<:AbstractVector{<:Real}`. Defaults to `Vector{Float64}`. Pass a different
 `S` (for example `Vector{Float32}` or `CuVector{Float64}`) to run AD in
 another precision or keep the tape on a GPU.
+
+The optional `coloring_algorithm` defaults to SMC's greedy coloring with
+substitution decompression.
 """
-struct Mode{S<:AbstractVector{<:Real}} <:
-       MOI.Nonlinear.AbstractAutomaticDifferentiation end
+struct Mode{S<:AbstractVector{<:Real},C<:SMC.GreedyColoringAlgorithm} <:
+       MOI.Nonlinear.AbstractAutomaticDifferentiation
+    coloring_algorithm::C
+end
+
+function Mode{S}(
+    coloring_algorithm::SMC.GreedyColoringAlgorithm = SMC.GreedyColoringAlgorithm(;
+        decompression = :substitution,
+    ),
+) where {S<:AbstractVector{<:Real}}
+    return Mode{S,typeof(coloring_algorithm)}(coloring_algorithm)
+end
+
+function Mode(coloring_algorithm::SMC.GreedyColoringAlgorithm)
+    return Mode{Vector{Float64}}(coloring_algorithm)
+end
 
 Mode() = Mode{Vector{Float64}}()
 
@@ -49,7 +67,7 @@ import NaNMath:
     pow,
     sqrt
 
-include("Coloring/Coloring.jl")
+include("coloring.jl")
 include("graph_tools.jl")
 include("sizes.jl")
 include("univariate_expressions.jl")
@@ -101,10 +119,17 @@ end
 # Create an ArrayDiff Evaluator from an ArrayDiff Model.
 function Evaluator(
     model::ArrayDiff.Model,
-    ::Mode{S},
+    mode::Mode{S},
     ordered_variables::Vector{MOI.VariableIndex},
 ) where {S<:AbstractVector{<:Real}}
-    return Evaluator(model, NLPEvaluator{eltype(S),S}(model, ordered_variables))
+    return Evaluator(
+        model,
+        NLPEvaluator{eltype(S),S}(
+            model,
+            ordered_variables,
+            mode.coloring_algorithm,
+        ),
+    )
 end
 
 # Called by solvers via MOI.Nonlinear.Evaluator(nlp_model, ad_backend, vars).