[Nonlinear] add support for atan2 (#1987)

Author: odow

src/FileFormats/NL/NLExpr.jl

@@ -62,7 +62,9 @@ const _JULIA_TO_AMPL = Dict{Symbol,Int}(
     :cosh => OP_cosh,
     :cos => OP_cos,
     :atanh => OP_atanh,
-    # OP_atan2 = 48,
+    # Because :atan is also a univariate operator, this is special cased in
+    # process_expr!.
+    # :atan = OP_atan2,
     :atan => OP_atan,
     :asinh => OP_asinh,
     :asin => OP_asin,
@@ -146,7 +148,7 @@ const _AMPL_TO_JULIA = Dict{Int,Tuple{Int,Function}}(
     OP_cosh => (1, cosh),
     OP_cos => (1, cos),
     OP_atanh => (1, atanh),
-    # OP_atan2 = 48,
+    OP_atan2 => (2, atan),
     OP_atan => (1, atan),
     OP_asinh => (1, asinh),
     OP_asin => (1, asin),
@@ -431,6 +433,11 @@ function _process_expr!(expr::_NLExpr, args::Vector{Any})
     if opcode === nothing
         error("Unsupported operation $(op)")
     end
+    if op == :atan && N == 2
+        # Special case binary use of atan, because Julia uses method overloading
+        # instead of having an explicit atan2 function.
+        opcode = OP_atan2
+    end
     push!(expr.nonlinear_terms, opcode)
     if opcode in _NARY_OPCODES
         push!(expr.nonlinear_terms, N)

src/Nonlinear/operators.jl

@@ -82,7 +82,7 @@ const DEFAULT_UNIVARIATE_OPERATORS = first.(SYMBOLIC_UNIVARIATE_EXPRESSIONS)
 
 The list of multivariate operators that are supported by default.
 """
-const DEFAULT_MULTIVARIATE_OPERATORS = [:+, :-, :*, :^, :/, :ifelse]
+const DEFAULT_MULTIVARIATE_OPERATORS = [:+, :-, :*, :^, :/, :ifelse, :atan]
 
 """
     OperatorRegistry()
@@ -544,6 +544,9 @@ function eval_multivariate_function(
     elseif op == :ifelse
         @assert length(x) == 3
         return ifelse(Bool(x[1]), x[2], x[3])
+    elseif op == :atan
+        @assert length(x) == 2
+        return atan(x[1], x[2])
     end
     id = registry.multivariate_operator_to_id[op]
     offset = id - registry.multivariate_user_operator_start
@@ -619,6 +622,11 @@ function eval_multivariate_gradient(
         g[1] = zero(T)  # It doesn't matter what this is.
         g[2] = x[1] == one(T)
         g[3] = x[1] == zero(T)
+    elseif op == :atan
+        @assert length(x) == 2
+        base = x[1]^2 + x[2]^2
+        g[1] = x[2] / base
+        g[2] = -x[1] / base
     else
         id = registry.multivariate_operator_to_id[op]
         offset = id - registry.multivariate_user_operator_start
@@ -707,6 +715,18 @@ function eval_multivariate_hessian(
         d = 1 / x[2]^2
         H[2, 1] = -d
         H[2, 2] = 2 * x[1] * d / x[2]
+    elseif op == :atan
+        # f(x)  = atan(y, x)
+        #
+        # ∇f(x) = +x/(x^2+y^2)
+        #         -y/(x^2+y^2)
+        #
+        # ∇²(x) = -(2xy)/(x^2+y^2)^2
+        #         (y^2-x^2)/(x^2+y^2)^2 (2xy)/(x^2+y^2)^2
+        base = (x[1]^2 + x[2]^2)^2
+        H[1, 1] = -2 * x[2] * x[1] / base
+        H[2, 1] = (x[1]^2 - x[2]^2) / base
+        H[2, 2] = 2 * x[2] * x[1] / base
     else
         id = registry.multivariate_operator_to_id[op]
         offset = id - registry.multivariate_user_operator_start

test/FileFormats/NL/NL.jl

@@ -212,6 +212,19 @@ function test_nlexpr_binary_specialcase()
     )
 end
 
+function test_nlexpr_atan_and_atan2()
+    x = MOI.VariableIndex(1)
+    y = MOI.VariableIndex(2)
+    _test_nlexpr(:(atan($x)), [NL.OP_atan, x], Dict(x => 0.0), 0.0)
+    _test_nlexpr(
+        :(atan($x, $y)),
+        [NL.OP_atan2, x, y],
+        Dict(x => 0.0, y => 0.0),
+        0.0,
+    )
+    return
+end
+
 function test_nlexpr_unsupportedoperation()
     x = MOI.VariableIndex(1)
     err = ErrorException("Unsupported operation foo")

test/FileFormats/NL/read.jl

@@ -118,6 +118,28 @@ function test_parse_header_binary()
     return
 end
 
+function test_parse_expr_atan2()
+    model = NL._CacheModel()
+    io = IOBuffer()
+    write(io, "o48\nv0\nv1\n")
+    seekstart(io)
+    x = MOI.VariableIndex.(1:2)
+    @test NL._parse_expr(io, model) == :(atan($(x[1]), $(x[2])))
+    @test eof(io)
+    return
+end
+
+function test_parse_expr_atan()
+    model = NL._CacheModel()
+    io = IOBuffer()
+    write(io, "o49\nv0\n")
+    seekstart(io)
+    x = MOI.VariableIndex.(1:1)
+    @test NL._parse_expr(io, model) == :(atan($(x[1])))
+    @test eof(io)
+    return
+end
+
 function test_parse_header_assertion_errors()
     model = NL._CacheModel()
     for header in [

test/Nonlinear/Nonlinear.jl

@@ -842,6 +842,34 @@ function test_NLPBlockData()
     return
 end
 
+function test_parse_atan2()
+    model = Nonlinear.Model()
+    x = MOI.VariableIndex(1)
+    y = MOI.VariableIndex(2)
+    θ = π / 4
+    Nonlinear.add_constraint(model, :(atan($x, $y)), MOI.LessThan(θ))
+    evaluator = Nonlinear.Evaluator(model)
+    MOI.initialize(evaluator, [:ExprGraph])
+    @test MOI.constraint_expr(evaluator, 1) == :(atan(x[$x], x[$y]) <= $θ)
+    return
+end
+
+function test_eval_atan2()
+    r = Nonlinear.OperatorRegistry()
+    x = [1.1, 2.2]
+    @test Nonlinear.eval_multivariate_function(r, :atan, x) ≈ atan(x[1], x[2])
+    g = zeros(2)
+    Nonlinear.eval_multivariate_gradient(r, :atan, g, x)
+    @test g[1] ≈ x[2] / (x[1]^2 + x[2]^2)
+    @test g[2] ≈ -x[1] / (x[1]^2 + x[2]^2)
+    H = LinearAlgebra.LowerTriangular(zeros(2, 2))
+    @test Nonlinear.eval_multivariate_hessian(r, :atan, H, x)
+    @test H[1, 1] ≈ -2 * x[2] * x[1] / (x[1]^2 + x[2]^2)^2
+    @test H[2, 1] ≈ (x[1]^2 - x[2]^2) / (x[1]^2 + x[2]^2)^2
+    @test H[2, 2] ≈ 2 * x[2] * x[1] / (x[1]^2 + x[2]^2)^2
+    return
+end
+
 end
 
 TestNonlinear.runtests()