Update special expression handling (#140)

src/EAGO.jl

@@ -30,6 +30,10 @@ module EAGO
     import Base: setindex!, + , *, -, ^, /, zero, one, inv, log, log10, exp, exp10, isempty, min, max
 
     @reexport using McCormick
+    import McCormick: relu, leaky_relu, maxsig, maxtanh, softplus, pentanh, sigmoid, bisigmoid,
+                      softsign, gelu, swish, xabsx, logcosh, xlogx, erf, erfinv, erfc, erfcinv, 
+                      xexpax, arh, param_relu, elu, selu, lower_bnd, upper_bnd, bnd, positive, negative
+
     @reexport using SpecialFunctions
     #@reexport using ReverseMcCormick
 
@@ -83,7 +87,9 @@ module EAGO
            EAGOModel, AuxiliaryVariableRef
     export auxiliary_variable, @auxiliary_variable, add_mimo_expression, @add_mimo_expression 
 
-    export register_eago_operators!
+    export relu, leaky_relu, maxsig, maxtanh, softplus, pentanh, sigmoid, bisigmoid,
+           softsign, gelu, swish, xabsx, logcosh, xlogx, erf, erfinv, erfc, erfcinv, 
+           xexpax, arh, param_relu, elu, selu, lower_bnd, upper_bnd, bnd, positive, negative
 
     # Map/reduce nonallocating no bounds checking map-reduce like utilities
     include(joinpath(@__DIR__, "eago_optimizer", "debug_tools.jl"))

src/eago_optimizer/functions/nonlinear/register_special.jl

@@ -7,55 +7,144 @@
 # https://github.com/PSORLab/EAGO.jl
 ################################################################################
 # src/eago_optimizer/functions/nonlinear/register_special.jl
-# Defines the function used to register nonstandard nonlinear terms.
+# Defines functions to register nonstandard nonlinear terms in JuMP models.
 ################################################################################
 
-"""
-register_eago_operators!
-
-Registers all nonstandard nonlinear terms available in EAGO in a JuMP. Uses of
-these is generally preferable in EAGO as the relaxations EAGO will generate
-will usually be tighter (speeding up convergence time). Note that this will
-work can be used by other nonlinear solvers (Ipopt for instance).
-"""
-function register_eago_operators!(m::JuMP.Model)
-
-    # Register activation functions w/o parameters
-    JuMP.register(m, :relu, 1, relu, McCormick.relu_deriv, McCormick.relu_deriv2)
-    JuMP.register(m, :leaky_relu, 1, leaky_relu, McCormick.leaky_relu_deriv, McCormick.leaky_relu_deriv2)
-    JuMP.register(m, :maxsig, 1, maxsig, McCormick.maxsig_deriv, McCormick.maxsig_deriv2)
-    JuMP.register(m, :maxtanh, 1, maxtanh, McCormick.maxtanh_deriv, McCormick.maxtanh_deriv2)
-    JuMP.register(m, :softplus, 1, softplus, McCormick.softplus_deriv, McCormick.softplus_deriv2)
-    JuMP.register(m, :pentanh, 1, pentanh, McCormick.pentanh_deriv, McCormick.pentanh_deriv2)
-    JuMP.register(m, :sigmoid, 1, sigmoid, McCormick.sigmoid_deriv, McCormick.sigmoid_deriv2)
-    JuMP.register(m, :bisigmoid, 1, bisigmoid, McCormick.bisigmoid_deriv, McCormick.bisigmoid_deriv2)
-    JuMP.register(m, :softsign, 1, softsign, McCormick.softsign_deriv, McCormick.softsign_deriv2)
-    JuMP.register(m, :gelu, 1, gelu, McCormick.gelu_deriv, McCormick.gelu_deriv2)
-    JuMP.register(m, :swish, 1, swish, McCormick.swish_deriv, McCormick.swish_deriv2)
-    JuMP.register(m, :xabsx, 1, xabsx, McCormick.xabsx_deriv, McCormick.xabsx_deriv2)
-    JuMP.register(m, :logcosh, 1, logcosh, McCormick.logcosh_deriv, McCormick.logcosh_deriv2)
-
-    # Register activation functions w/ parameters
-    MOINL.register_operator(m.nlp_model, :param_relu, 2, param_relu, McCormick.param_relu_grad)
-    MOINL.register_operator(m.nlp_model, :elu, 2, elu, McCormick.elu_grad)
-    MOINL.register_operator(m.nlp_model, :selu, 3, selu, McCormick.selu_grad)
-
-    # Register other functions
-    JuMP.register(m, :xlogx, 1, xlogx, McCormick.xlogx_deriv, McCormick.xlogx_deriv2)
-    JuMP.register(m, :f_erf, 1, x -> erf(x), McCormick.erf_deriv, McCormick.erf_deriv2)
-    JuMP.register(m, :f_erfinv, 1, x -> erfinv(x), McCormick.erfinv_deriv, McCormick.erfinv_deriv2)
-    JuMP.register(m, :f_erfc, 1, x -> erfc(x), McCormick.erfc_deriv, McCormick.erfc_deriv2)
-    JuMP.register(m, :f_erfcinv, 1, x -> erfcinv(x), x -> -McCormick.erfinv_deriv(1.0 - x), x -> McCormick.erfinv_deriv2(1.0 - x))
-
-    MOINL.register_operator(m.nlp_model, :arh, 2, arh, McCormick.arh_grad)
-    MOINL.register_operator(m.nlp_model, :xexpax, 2, xexpax, McCormick.xexpax_grad)
-
-    # Register bounding functions
-    MOINL.register_operator(m.nlp_model, :lower_bnd, 2, lower_bnd, McCormick.d_lower_bnd_grad)
-    MOINL.register_operator(m.nlp_model, :upper_bnd, 2, upper_bnd, McCormick.d_upper_bnd_grad)
-    MOINL.register_operator(m.nlp_model, :bnd, 3, bnd, McCormick.d_bnd_grad)
-    JuMP.register(m, :positive, 1, positive, x -> 1.0, x -> 0.0)
-    JuMP.register(m, :negative, 1, negative, x -> 1.0, x -> 0.0)
-
-    return nothing
+# Operators with 1 input
+for operator in (:relu, :leaky_relu, :maxsig, :maxtanh, :softplus, :pentanh, :sigmoid, :bisigmoid,
+                 :softsign, :gelu, :swish, :xabsx, :logcosh, :xlogx, :erf, :erfinv, :erfc)
+    func = quote
+        model = owner_model(x)
+        op = get(model.obj_dict, Symbol("EAGO_", $(String(operator))), nothing)
+        if op === nothing
+            op = add_nonlinear_operator(
+                model,
+                1,
+                getfield(McCormick, Symbol($(String(operator)))),
+                getfield(McCormick, Symbol($(String(operator)), "_deriv")),
+                getfield(McCormick, Symbol($(String(operator)), "_deriv2"));
+                name = Symbol("EAGO_", $(String(operator))),
+            )
+            model[Symbol("EAGO_", $(String(operator)))] = op
+        end
+        return op(x)
+    end
+
+    @eval @inline ($operator)(x::VariableRef) = $func
+end
+
+# Operator for `McCormick.erfcinv`
+function erfcinv(x::VariableRef)
+    model = owner_model(x)
+    op = get(model.obj_dict, :EAGO_erfcinv, nothing)
+    if op === nothing
+        op = add_nonlinear_operator(
+            model,
+            1,
+            McCormick.erfcinv,
+            x -> -McCormick.erfinv_deriv(1.0 - x),
+            x -> McCormick.erfinv_deriv2(1.0 - x);
+            name = :EAGO_erfcinv,
+        )
+        model[:EAGO_erfcinv] = op
+    end
+    return op(x)
+end
+
+# Operators with 2 inputs
+for operator in (:xexpax, :arh, :param_relu, :elu)
+    func = quote
+        model = owner_model(x)
+        op = get(model.obj_dict, Symbol("EAGO_", $operator), nothing)
+        if op === nothing
+            op = add_nonlinear_operator(
+                model,
+                2,
+                getfield(McCormick, Symbol($operator)),
+                getfield(McCormick, Symbol($operator, "_grad"));
+                name = Symbol("EAGO_", $operator),
+            )
+            model[Symbol("EAGO_", $operator)] = op
+        end
+        return op(x, p)
+    end
+
+    @eval @inline ($operator)(x::VariableRef, p::VariableRef) = $func
+end
+
+# Operator for `McCormick.selu`
+function selu(x::VariableRef, α::VariableRef, λ::VariableRef)
+    model = owner_model(x)
+    op = get(model.obj_dict, :EAGO_selu, nothing)
+    if op === nothing
+        op = add_nonlinear_operator(
+            model,
+            3,
+            McCormick.selu,
+            McCormick.selu_grad;
+            name = :EAGO_selu,
+        )
+        model[:EAGO_selu] = op
+    end
+    return op(x, α, λ)
+end
+
+# Operators for `McCormick.lower_bnd` and `McCormick.upper_bnd`
+for operator in (:lower_bnd, :upper_bnd)
+    func = quote
+        model = owner_model(x)
+        op = get(model.obj_dict, Symbol("EAGO_", $operator), nothing)
+        if op === nothing
+            op = add_nonlinear_operator(
+                model,
+                2,
+                getfield(McCormick, Symbol($operator)),
+                getfield(McCormick, Symbol("d_", $operator, "_grad"));
+                name = Symbol("EAGO_", $operator),
+            )
+            model[Symbol("EAGO_", $operator)] = op
+        end
+        return op(x, b)
+    end
+
+    @eval @inline ($operator)(x::VariableRef, b::VariableRef) = $func
+end
+
+# Operator for `McCormick.bnd`
+function bnd(x::VariableRef, lb::VariableRef, ub::VariableRef)
+    model = owner_model(x)
+    op = get(model.obj_dict, :EAGO_bnd, nothing)
+    if op === nothing
+        op = add_nonlinear_operator(
+            model,
+            3,
+            McCormick.bnd,
+            McCormick.d_bnd_grad;
+            name = :EAGO_bnd,
+        )
+        model[:EAGO_bnd] = op
+    end
+    return op(x, lb, ub)
+end
+
+# Operators for `McCormick.positive` and `McCormick.negative`
+for operator in (:positive, :negative)
+    func = quote
+        model = owner_model(x)
+        op = get(model.obj_dict, Symbol("EAGO_", $operator), nothing)
+        if op === nothing
+            op = add_nonlinear_operator(
+                model,
+                1,
+                getfield(McCormick, Symbol($operator)),
+                x -> 1.0,
+                x -> 0.0;
+                name = Symbol("EAGO_", $operator),
+            )
+        model[Symbol("EAGO_", $operator)] = op
+        end
+        return op(x)
+    end
+
+    @eval @inline ($operator)(x::VariableRef) = $func
 end

src/eago_optimizer/moi_wrapper.jl

@@ -279,12 +279,15 @@ function MOI.set(m::Optimizer, ::MOI.ObjectiveFunction{T}, f::T) where T <: Unio
     if f isa MOI.ScalarNonlinearFunction
         if isnothing(m._input_problem._nlp_data)
             model = MOI.Nonlinear.Model()
+            MOI.Nonlinear.set_objective(model, f)
             backend = MOI.Nonlinear.SparseReverseMode()
             vars = MOI.get(m, MOI.ListOfVariableIndices())
             evaluator = MOI.Nonlinear.Evaluator(model, backend, vars) 
             m._input_problem._nlp_data = MOI.NLPBlockData(evaluator)
+        else
+            MOI.Nonlinear.set_objective(m._input_problem._nlp_data.evaluator.model, f)
+            m._input_problem._nlp_data = MOI.NLPBlockData(m._input_problem._nlp_data.evaluator)
         end
-        MOI.Nonlinear.set_objective(m._input_problem._nlp_data.evaluator.model, f)
     else
         m._input_problem._objective = f
     end