From 24012962dd89e36b413469a3c85b26491898c749 Mon Sep 17 00:00:00 2001
From: Oscar Dowson <odow@users.noreply.github.com>
Date: Tue, 10 Sep 2024 09:29:56 +1200
Subject: [PATCH] [Nonlinear] allow univariate operators with only gradient
 information (#2542)

---
 .../ReverseAD/mathoptinterface_api.jl         | 14 +++++++------
 src/Nonlinear/operators.jl                    | 20 ++++++++++++++-----
 test/Nonlinear/ReverseAD.jl                   | 17 ++++++++++++++++
 3 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl
index f7fcd9615d..eec1cf1fb3 100644
--- a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl
+++ b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl
@@ -4,13 +4,15 @@
 # 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.
 
+_no_hessian(op::MOI.Nonlinear._UnivariateOperator) = op.f′′ === nothing
+_no_hessian(op::MOI.Nonlinear._MultivariateOperator) = op.∇²f === nothing
+
 function MOI.features_available(d::NLPEvaluator)
-    # Check if we are missing any hessians for user-defined multivariate
-    # operators, in which case we need to disable :Hess and :HessVec.
-    d.disable_2ndorder = any(
-        op -> op.∇²f === nothing,
-        d.data.operators.registered_multivariate_operators,
-    )
+    # Check if we are missing any hessians for user-defined operators, in which
+    # case we need to disable :Hess and :HessVec.
+    d.disable_2ndorder =
+        any(_no_hessian, d.data.operators.registered_univariate_operators) ||
+        any(_no_hessian, d.data.operators.registered_multivariate_operators)
     if d.disable_2ndorder
         return [:Grad, :Jac, :JacVec]
     end
diff --git a/src/Nonlinear/operators.jl b/src/Nonlinear/operators.jl
index e4f51e3c5b..e0a94f414c 100644
--- a/src/Nonlinear/operators.jl
+++ b/src/Nonlinear/operators.jl
@@ -54,6 +54,13 @@ struct _UnivariateOperator{F,F′,F′′}
     f::F
     f′::F′
     f′′::F′′
+    function _UnivariateOperator(
+        f::Function,
+        f′::Function,
+        f′′::Union{Nothing,Function} = nothing,
+    )
+        return new{typeof(f),typeof(f′),typeof(f′′)}(f, f′, f′′)
+    end
 end
 
 struct _MultivariateOperator{F,F′,F′′}
@@ -339,14 +346,17 @@ end
 function _UnivariateOperator(op::Symbol, f::Function)
     _validate_register_assumptions(f, op, 1)
     f′ = _checked_derivative(f, op)
-    f′′ = _checked_derivative(f′, op)
-    return _UnivariateOperator(f, f′, f′′)
+    return _UnivariateOperator(op, f, f′)
 end
 
 function _UnivariateOperator(op::Symbol, f::Function, f′::Function)
-    _validate_register_assumptions(f′, op, 1)
-    f′′ = _checked_derivative(f′, op)
-    return _UnivariateOperator(f, f′, f′′)
+    try
+        _validate_register_assumptions(f′, op, 1)
+        f′′ = _checked_derivative(f′, op)
+        return _UnivariateOperator(f, f′, f′′)
+    catch
+        return _UnivariateOperator(f, f′, nothing)
+    end
 end
 
 function _UnivariateOperator(::Symbol, f::Function, f′::Function, f′′::Function)
diff --git a/test/Nonlinear/ReverseAD.jl b/test/Nonlinear/ReverseAD.jl
index 8eb597f873..60694e7671 100644
--- a/test/Nonlinear/ReverseAD.jl
+++ b/test/Nonlinear/ReverseAD.jl
@@ -1135,6 +1135,23 @@ function test_varying_length_x()
     return
 end
 
+function test_univariate_operator_with_no_second_order()
+    f(x::Float64) = x^2
+    df(x::Float64) = 2 * x
+    model = MOI.Nonlinear.Model()
+    MOI.Nonlinear.register_operator(model, :op_f, 1, f, df)
+    x = MOI.VariableIndex(1)
+    MOI.Nonlinear.add_constraint(model, :(op_f($x)), MOI.LessThan(2.0))
+    evaluator =
+        MOI.Nonlinear.Evaluator(model, MOI.Nonlinear.SparseReverseMode(), [x])
+    @test !(:Hess in MOI.features_available(evaluator))
+    MOI.initialize(evaluator, [:Grad, :Jac])
+    J = zeros(length(MOI.jacobian_structure(evaluator)))
+    MOI.eval_constraint_jacobian(evaluator, J, [2.0])
+    @test J == [4.0]
+    return
+end
+
 end  # module
 
 TestReverseAD.runtests()