Use an enum for the objective sink of ModelWithQuad

Author: blegat

src/Nonlinear/model_with_quad.jl

@@ -4,11 +4,14 @@
 # Use of this source code is governed by an MIT-style license that can be found
 # in the LICENSE.md file or at https://opensource.org/licenses/MIT.
 
+# Where the objective of a `ModelWithQuad` currently lives.
+@enum(_ObjectiveSink, _NONE, _QUAD, _INNER)
+
 """
     ModelWithQuad{T,M}(
         qp::QPBlockData{T},
         inner::M;
-        objective_sink::Symbol = :none,
+        objective_sink::_ObjectiveSink = _NONE,
     ) where {T,M}
 
 A model layer that owns the variables of the model, stores affine and
@@ -29,8 +32,8 @@ that storage as `parameters::Vector{T}`, like [`Model`](@ref) does:
 Add constraints with [`add_constraint`](@ref) or `MOI.add_constraint`, and
 set the objective with [`set_objective`](@ref): affine and quadratic
 functions are routed to the QP block, everything else to the inner model.
-`objective_sink` records where the objective currently lives (`:none`,
-`:quad` or `:inner`).
+`objective_sink` records where the objective currently lives (`_NONE`,
+`_QUAD` or `_INNER`).
 
 Create the corresponding evaluator, [`EvaluatorWithQuad`](@ref), with
 `Evaluator(model, backend)`, or construct it directly from an inner
@@ -41,12 +44,12 @@ mutable struct ModelWithQuad{T,M}
     variables::MOI.Utilities.VariablesContainer{T}
     qp::QPBlockData{T}
     inner::M
-    objective_sink::Symbol # :none, :quad or :inner
+    objective_sink::_ObjectiveSink
 
     function ModelWithQuad{T}(
         qp::QPBlockData{T},
         inner::M;
-        objective_sink::Symbol = :none,
+        objective_sink::_ObjectiveSink = _NONE,
     ) where {T,M}
         model = new{T,M}(
             MOI.Utilities.VariablesContainer{T}(),
@@ -223,7 +226,7 @@ function set_objective(
 ) where {T}
     MOI.set(model.qp, MOI.ObjectiveFunction{typeof(obj)}(), obj)
     set_objective(model.inner, nothing)
-    model.objective_sink = :quad
+    model.objective_sink = _QUAD
     return
 end
 
@@ -234,7 +237,7 @@ function set_objective(model::ModelWithQuad{T}, obj) where {T}
         obj = _replace_parameters(obj)
     end
     set_objective(model.inner, obj)
-    model.objective_sink = obj === nothing ? :none : :inner
+    model.objective_sink = obj === nothing ? _NONE : _INNER
     return
 end
 
@@ -385,9 +388,9 @@ end
 
 function MOI.eval_objective(d::EvaluatorWithQuad{T}, x) where {T}
     sink = d.model.objective_sink
-    if sink == :quad
+    if sink == _QUAD
         return MOI.eval_objective(d.model.qp, x)
-    elseif sink == :inner
+    elseif sink == _INNER
         return MOI.eval_objective(d.inner, x)
     else
         return zero(T)
@@ -396,9 +399,9 @@ end
 
 function MOI.eval_objective_gradient(d::EvaluatorWithQuad{T}, grad, x) where {T}
     sink = d.model.objective_sink
-    if sink == :quad
+    if sink == _QUAD
         MOI.eval_objective_gradient(d.model.qp, grad, x)
-    elseif sink == :inner
+    elseif sink == _INNER
         MOI.eval_objective_gradient(d.inner, grad, x)
     else
         grad .= zero(T)
@@ -534,7 +537,7 @@ end
 _has_objective(d::Evaluator) = d.model.objective !== nothing
 
 function _has_objective(d::EvaluatorWithQuad)
-    if d.model.objective_sink == :quad
+    if d.model.objective_sink == _QUAD
         return true
     end
     return _has_objective(d.inner)

test/Nonlinear/test_model_with_quad.jl

@@ -165,14 +165,14 @@ end
 function test_objective_sink_switching()
     model = Nonlinear.ModelWithQuad(Nonlinear.Model())
     x = MOI.add_variable(model)
-    @test model.objective_sink == :none
+    @test model.objective_sink == Nonlinear._NONE
     f = MOI.ScalarQuadraticFunction(
         [MOI.ScalarQuadraticTerm(2.0, x, x)],
         MOI.ScalarAffineTerm{Float64}[],
         0.0,
     )
     Nonlinear.set_objective(model, f)
-    @test model.objective_sink == :quad
+    @test model.objective_sink == Nonlinear._QUAD
     @test MOI.get(model, MOI.ObjectiveFunctionType()) ==
           MOI.ScalarQuadraticFunction{Float64}
     @test MOI.get(model, MOI.ObjectiveFunction{typeof(f)}()) ≈ f
@@ -183,7 +183,7 @@ function test_objective_sink_switching()
     # Switch to a nonlinear objective: the quadratic objective must be
     # cleared, including its Hessian entries.
     Nonlinear.set_objective(model, :(sin($x)))
-    @test model.objective_sink == :inner
+    @test model.objective_sink == Nonlinear._INNER
     d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode())
     MOI.initialize(d, [:Grad, :Jac, :Hess])
     @test MOI.eval_objective(d, [3.0]) == sin(3.0)
@@ -195,12 +195,12 @@ function test_objective_sink_switching()
     # Switch to a linear objective, and then remove it.
     g = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(2.0, x)], 1.0)
     Nonlinear.set_objective(model, g)
-    @test model.objective_sink == :quad
+    @test model.objective_sink == Nonlinear._QUAD
     d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode())
     MOI.initialize(d, [:Grad, :Jac])
     @test MOI.eval_objective(d, [3.0]) == 7.0
     Nonlinear.set_objective(model, nothing)
-    @test model.objective_sink == :none
+    @test model.objective_sink == Nonlinear._NONE
     d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode())
     MOI.initialize(d, [:Grad, :Jac])
     @test MOI.eval_objective(d, [3.0]) == 0.0