Add ScalarNonlinearFunction (#2059)

Author: odow

docs/src/manual/standard_form.md

@@ -37,6 +37,7 @@ The function types implemented in MathOptInterface.jl are:
 | [`VariableIndex`](@ref) | ``x_j``, the projection onto a single coordinate defined by a variable index ``j``. |
 | [`VectorOfVariables`](@ref) | The projection onto multiple coordinates (that is, extracting a sub-vector). |
 | [`ScalarAffineFunction`](@ref) | ``a^T x + b``, where ``a`` is a vector and ``b`` scalar. |
+| [`ScalarNonlinearFunction`](@ref) | ``f(x)``, where ``f`` is a nonlinear function. |
 | [`VectorAffineFunction`](@ref) | ``A x + b``, where ``A`` is a matrix and ``b`` is a vector. |
 | [`ScalarQuadraticFunction`](@ref) | ``\frac{1}{2} x^T Q x + a^T x + b``, where ``Q`` is a symmetric matrix, ``a`` is a vector, and ``b`` is a constant. |
 | [`VectorQuadraticFunction`](@ref) | A vector of scalar-valued quadratic functions. |

docs/src/reference/errors.md

@@ -75,6 +75,7 @@ ModifyObjectiveNotAllowed
 DeleteNotAllowed
 UnsupportedSubmittable
 SubmitNotAllowed
+UnsupportedNonlinearOperator
 ```
 
 Note that setting the [`ConstraintFunction`](@ref) of a [`VariableIndex`](@ref)

docs/src/reference/models.md

@@ -50,6 +50,8 @@ ListOfOptimizerAttributesSet
 ListOfModelAttributesSet
 ListOfVariableAttributesSet
 ListOfConstraintAttributesSet
+UserDefinedFunction
+ListOfSupportedNonlinearOperators
 ```
 
 ## Optimizer interface

docs/src/reference/standard_form.md

@@ -25,6 +25,7 @@ ScalarAffineTerm
 ScalarAffineFunction
 ScalarQuadraticTerm
 ScalarQuadraticFunction
+ScalarNonlinearFunction
 ```
 
 ## Vector functions

src/Bridges/Objective/bridges/slack.jl

@@ -70,7 +70,7 @@ function bridge_objective(
     end
     constraint = MOI.Utilities.normalize_and_add_constraint(model, f, set)
     MOI.set(model, MOI.ObjectiveFunction{MOI.VariableIndex}(), slack)
-    return SlackBridge{T,F,G}(slack, constraint, MOI.constant(f))
+    return SlackBridge{T,F,G}(slack, constraint, MOI.constant(f, T))
 end
 
 function supports_objective_function(
@@ -166,7 +166,11 @@ function MOI.get(
     bridge::SlackBridge{T,F,G},
 ) where {T,F,G<:MOI.AbstractScalarFunction}
     func = MOI.get(model, MOI.ConstraintFunction(), bridge.constraint)
-    f = MOI.Utilities.operate(+, T, func, bridge.constant)
+    f = if !iszero(bridge.constant)
+        MOI.Utilities.operate(+, T, func, bridge.constant)
+    else
+        func
+    end
     g = MOI.Utilities.remove_variable(f, bridge.slack)
     return MOI.Utilities.convert_approx(G, g)
 end

src/Nonlinear/model.jl

@@ -327,3 +327,7 @@ function evaluate(
     end
     return storage[1]
 end
+
+function MOI.get(model::Model, attr::MOI.ListOfSupportedNonlinearOperators)
+    return MOI.get(model.operators, attr)
+end

src/Nonlinear/operators.jl

@@ -74,13 +74,60 @@ end
     DEFAULT_UNIVARIATE_OPERATORS
 
 The list of univariate operators that are supported by default.
+
+## Example
+
+```jldoctest
+julia> import MathOptInterface as MOI
+
+julia> MOI.Nonlinear.DEFAULT_UNIVARIATE_OPERATORS
+72-element Vector{Symbol}:
+ :+
+ :-
+ :abs
+ :sqrt
+ :cbrt
+ :abs2
+ :inv
+ :log
+ :log10
+ :log2
+ ⋮
+ :airybi
+ :airyaiprime
+ :airybiprime
+ :besselj0
+ :besselj1
+ :bessely0
+ :bessely1
+ :erfcx
+ :dawson
+```
 """
 const DEFAULT_UNIVARIATE_OPERATORS = first.(SYMBOLIC_UNIVARIATE_EXPRESSIONS)
 
 """
     DEFAULT_MULTIVARIATE_OPERATORS
 
 The list of multivariate operators that are supported by default.
+
+## Example
+
+```jldoctest
+julia> import MathOptInterface as MOI
+
+julia> MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS
+9-element Vector{Symbol}:
+ :+
+ :-
+ :*
+ :^
+ :/
+ :ifelse
+ :atan
+ :min
+ :max
+```
 """
 const DEFAULT_MULTIVARIATE_OPERATORS =
     [:+, :-, :*, :^, :/, :ifelse, :atan, :min, :max]
@@ -140,6 +187,19 @@ struct OperatorRegistry
     end
 end
 
+function MOI.get(
+    registry::OperatorRegistry,
+    ::MOI.ListOfSupportedNonlinearOperators,
+)
+    ops = vcat(
+        registry.univariate_operators,
+        registry.multivariate_operators,
+        registry.logic_operators,
+        registry.comparison_operators,
+    )
+    return unique(ops)
+end
+
 const _FORWARD_DIFF_METHOD_ERROR_HELPER = raw"""
 Common reasons for this include:
 

src/Nonlinear/parse.jl

@@ -36,6 +36,58 @@ function parse_expression(::Model, ::Expression, x::Any, ::Int)
     )
 end
 
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::MOI.ScalarNonlinearFunction,
+    parent_index::Int,
+)
+    stack = Tuple{Int,Any}[(parent_index, x)]
+    while !isempty(stack)
+        parent_node, arg = pop!(stack)
+        if arg isa MOI.ScalarNonlinearFunction
+            _parse_without_recursion_inner(stack, data, expr, arg, parent_node)
+        else
+            # We can use recursion here, because ScalarNonlinearFunction only
+            # occur in other ScalarNonlinearFunction.
+            parse_expression(data, expr, arg, parent_node)
+        end
+    end
+    return
+end
+
+function _get_node_type(data, x)
+    id = get(data.operators.univariate_operator_to_id, x.head, nothing)
+    if length(x.args) == 1 && id !== nothing
+        return id, MOI.Nonlinear.NODE_CALL_UNIVARIATE
+    end
+    id = get(data.operators.multivariate_operator_to_id, x.head, nothing)
+    if id !== nothing
+        return id, MOI.Nonlinear.NODE_CALL_MULTIVARIATE
+    end
+    id = get(data.operators.comparison_operator_to_id, x.head, nothing)
+    if id !== nothing
+        return id, MOI.Nonlinear.NODE_COMPARISON
+    end
+    id = get(data.operators.logic_operator_to_id, x.head, nothing)
+    if id !== nothing
+        return id, MOI.Nonlinear.NODE_LOGIC
+    end
+    return throw(MOI.UnsupportedNonlinearOperator(x.head))
+end
+
+function _parse_without_recursion_inner(stack, data, expr, x, parent)
+    id, node_type = _get_node_type(data, x)
+    push!(expr.nodes, Node(node_type, id, parent))
+    parent = length(expr.nodes)
+    # Args need to be pushed onto the stack in reverse because the stack is a
+    # first-in last-out datastructure.
+    for arg in reverse(x.args)
+        push!(stack, (parent, arg))
+    end
+    return
+end
+
 function parse_expression(
     data::Model,
     expr::Expression,
@@ -108,7 +160,7 @@ function _parse_univariate_expression(
             _parse_multivariate_expression(stack, data, expr, x, parent_index)
             return
         end
-        error("Unable to parse: $x")
+        throw(MOI.UnsupportedNonlinearOperator(x.args[1]))
     end
     push!(expr.nodes, Node(NODE_CALL_UNIVARIATE, id, parent_index))
     push!(stack, (length(expr.nodes), x.args[2]))
@@ -200,6 +252,28 @@ function parse_expression(
     return
 end
 
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::MOI.ScalarAffineFunction,
+    parent_index::Int,
+)
+    f = convert(MOI.ScalarNonlinearFunction, x)
+    parse_expression(data, expr, f, parent_index)
+    return
+end
+
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::MOI.ScalarQuadraticFunction,
+    parent_index::Int,
+)
+    f = convert(MOI.ScalarNonlinearFunction, x)
+    parse_expression(data, expr, f, parent_index)
+    return
+end
+
 function parse_expression(::Model, expr::Expression, x::Real, parent_index::Int)
     push!(expr.values, convert(Float64, x)::Float64)
     push!(expr.nodes, Node(NODE_VALUE, length(expr.values), parent_index))

src/Test/test_basic_constraint.jl

@@ -66,6 +66,17 @@ function _function(
     )
 end
 
+function _function(
+    ::Type{T},
+    ::Type{MOI.ScalarNonlinearFunction},
+    x::Vector{MOI.VariableIndex},
+) where {T}
+    return MOI.ScalarNonlinearFunction(
+        :+,
+        Any[MOI.ScalarNonlinearFunction(:^, Any[xi, 2]) for xi in x],
+    )
+end
+
 # Default fallback.
 _set(::Any, ::Type{S}) where {S} = _set(S)
 
@@ -316,7 +327,12 @@ for s in [
 ]
     S = getfield(MOI, s)
     functions = if S <: MOI.AbstractScalarSet
-        (:VariableIndex, :ScalarAffineFunction, :ScalarQuadraticFunction)
+        (
+            :VariableIndex,
+            :ScalarAffineFunction,
+            :ScalarQuadraticFunction,
+            :ScalarNonlinearFunction,
+        )
     else
         (:VectorOfVariables, :VectorAffineFunction, :VectorQuadraticFunction)
     end