Add constraint_gradient_structure and eval_constraint_gradient (#2200)

Author: odow

docs/src/reference/nonlinear.md

@@ -32,6 +32,8 @@ eval_objective
 eval_constraint
 eval_objective_gradient
 jacobian_structure
+eval_constraint_gradient
+constraint_gradient_structure
 eval_constraint_jacobian
 eval_constraint_jacobian_product
 eval_constraint_jacobian_transpose_product

src/Nonlinear/ReverseAD/mathoptinterface_api.jl

@@ -207,6 +207,22 @@ function MOI.jacobian_structure(d::NLPEvaluator)
     return J
 end
 
+function MOI.constraint_gradient_structure(d::NLPEvaluator, i)
+    return copy(d.constraints[i].grad_sparsity)
+end
+
+function MOI.eval_constraint_gradient(d::NLPEvaluator, ∇g, x, i)
+    _reverse_mode(d, x)
+    for k in d.constraints[i].grad_sparsity
+        d.jac_storage[k] = 0.0
+    end
+    _extract_reverse_pass(d.jac_storage, d, d.constraints[i])
+    for (j, k) in enumerate(d.constraints[i].grad_sparsity)
+        ∇g[j] = d.jac_storage[k]
+    end
+    return
+end
+
 function MOI.eval_constraint_jacobian(d::NLPEvaluator, J, x)
     _reverse_mode(d, x)
     fill!(J, 0.0)

src/Nonlinear/evaluator.jl

@@ -68,6 +68,7 @@ function MOI.initialize(evaluator::Evaluator, features::Vector{Symbol})
     evaluator.eval_objective_timer = 0.0
     evaluator.eval_objective_gradient_timer = 0.0
     evaluator.eval_constraint_timer = 0.0
+    evaluator.eval_constraint_gradient_timer = 0.0
     evaluator.eval_constraint_jacobian_timer = 0.0
     evaluator.eval_hessian_objective_timer = 0.0
     evaluator.eval_hessian_constraint_timer = 0.0
@@ -134,6 +135,17 @@ function MOI.eval_constraint(evaluator::Evaluator, g, x)
     return
 end
 
+function MOI.eval_constraint_gradient(evaluator::Evaluator, ∇g, x, i)
+    start = time()
+    MOI.eval_constraint_gradient(evaluator.backend, ∇g, x, i)
+    evaluator.eval_constraint_gradient_timer += time() - start
+    return
+end
+
+function MOI.constraint_gradient_structure(evaluator::Evaluator, i)
+    return MOI.constraint_gradient_structure(evaluator.backend, i)
+end
+
 function MOI.jacobian_structure(evaluator::Evaluator)
     return MOI.jacobian_structure(evaluator.backend)
 end

src/Nonlinear/types.jl

@@ -208,6 +208,7 @@ mutable struct Evaluator{B} <: MOI.AbstractNLPEvaluator
     eval_objective_timer::Float64
     eval_constraint_timer::Float64
     eval_objective_gradient_timer::Float64
+    eval_constraint_gradient_timer::Float64
     eval_constraint_jacobian_timer::Float64
     eval_hessian_objective_timer::Float64
     eval_hessian_constraint_timer::Float64
@@ -229,6 +230,7 @@ mutable struct Evaluator{B} <: MOI.AbstractNLPEvaluator
             0.0,
             0.0,
             0.0,
+            0.0,
         )
     end
 end

src/nlp.jl

@@ -264,6 +264,40 @@ The sparsity structure is assumed to be independent of the point ``x``.
 """
 function hessian_lagrangian_structure end
 
+"""
+    constraint_gradient_structure(d::AbstractNLPEvaluator, i::Int)::Vector{Int64}
+
+Returns a vector of indices, where each element indicates the position of a
+structurally nonzero element in the gradient of constraint ``\\nabla g_i(x)``.
+
+The indices are not required to be sorted and can contain duplicates, in which
+case the solver should combine the corresponding elements by adding them
+together.
+
+The sparsity structure is assumed to be independent of the point ``x``.
+"""
+function constraint_gradient_structure end
+
+"""
+    eval_constraint_gradient(
+        d::AbstractNLPEvaluator,
+        ∇g::AbstractVector{T},
+        x::AbstractVector{T},
+        i::Int,
+    )::Nothing where {T}
+
+Evaluate the gradient of constraint `i`, ``\\nabla g_i(x)``, and store the
+non-zero values in `∇g`, corresponding to the structure returned by
+[`constraint_gradient_structure`](@ref).
+
+## Implementation notes
+
+When implementing this method, you must not assume that `∇g` is
+`Vector{Float64}`, but you may assume that it supports `setindex!` and `length`.
+For example, it may be the `view` of a vector.
+"""
+function eval_constraint_gradient end
+
 """
     eval_constraint_jacobian(d::AbstractNLPEvaluator,
         J::AbstractVector{T},

test/Nonlinear/ReverseAD.jl

@@ -1029,6 +1029,28 @@ function test_pow_complex_result()
     return
 end
 
+function test_constraint_gradient()
+    x = MOI.VariableIndex(1)
+    y = MOI.VariableIndex(2)
+    model = Nonlinear.Model()
+    Nonlinear.add_constraint(model, :($x^2 + $x * $y + $y^2), MOI.LessThan(2.0))
+    Nonlinear.add_constraint(model, :(cos($y)), MOI.LessThan(2.0))
+    evaluator =
+        Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x, y])
+    MOI.initialize(evaluator, [:Grad, :Jac])
+    @test MOI.constraint_gradient_structure(evaluator, 1) == [1, 2]
+    @test MOI.constraint_gradient_structure(evaluator, 2) == [2]
+    x_val = [1.2, 2.3]
+    ∇g = [NaN, NaN]
+    MOI.eval_constraint_gradient(evaluator, ∇g, x_val, 1)
+    @test ∇g ≈ [2 * x_val[1] + x_val[2], x_val[1] + 2 * x_val[2]]
+    x_val = [1.2, 2.3]
+    ∇g = [NaN]
+    MOI.eval_constraint_gradient(evaluator, ∇g, x_val, 2)
+    @test ∇g ≈ [-sin(x_val[2])]
+    return
+end
+
 end  # module
 
 TestReverseAD.runtests()