Add @force_nonlinear macro for modifying how expressions are parsed (#…

…3732)

docs/src/manual/nonlinear.md

@@ -355,6 +355,34 @@ julia> expr.args
  x
 ```
 
+### Forcing nonlinear expressions
+
+The JuMP macros and operator overloading will preferentially build affine ([`GenericAffExpr`](@ref)) and quadratic ([`GenericQuadExpr`](@ref)) expressions
+instead of [`GenericNonlinearExpr`](@ref). For example:
+```jldoctest force_nonlinear
+julia> model = Model();
+
+julia> @variable(model, x);
+
+julia> f = (x - 0.1)^2
+x² - 0.2 x + 0.010000000000000002
+
+julia> typeof(f)
+QuadExpr (alias for GenericQuadExpr{Float64, GenericVariableRef{Float64}})
+```
+To override this behavior, use the [`@force_nonlinear`](@ref) macro:
+```jldoctest force_nonlinear
+julia> g = @force_nonlinear((x - 0.1)^2)
+(x - 0.1) ^ 2
+
+julia> typeof(g)
+NonlinearExpr (alias for GenericNonlinearExpr{GenericVariableRef{Float64}})
+```
+
+!!! warning
+    Use this macro only if necessary. See the docstring of [`@force_nonlinear`](@ref)
+    for more details on when you should use it.
+
 ## Function tracing
 
 Nonlinear expressions can be constructed using _function tracing_. Function

src/macros.jl

@@ -467,8 +467,6 @@ function _plural_macro_code(model, block, macro_sym)
     return code
 end
 
-include("macros/@objective.jl")
-include("macros/@expression.jl")
-include("macros/@constraint.jl")
-include("macros/@variable.jl")
-include("macros/@NL.jl")
+for file in readdir(joinpath(@__DIR__, "macros"))
+    include(joinpath(@__DIR__, "macros", file))
+end

src/macros/@force_nonlinear.jl

@@ -0,0 +1,121 @@
+#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+const _op_add = NonlinearOperator(+, :+)
+const _op_sub = NonlinearOperator(-, :-)
+const _op_mul = NonlinearOperator(*, :*)
+const _op_div = NonlinearOperator(/, :/)
+const _op_pow = NonlinearOperator(^, :^)
+
+"""
+    @force_nonlinear(expr)
+
+Change the parsing of `expr` to construct [`GenericNonlinearExpr`](@ref) instead
+of [`GenericAffExpr`](@ref) or [`GenericQuadExpr`](@ref).
+
+This macro works by walking `expr` and substituting all calls to `+`, `-`, `*`,
+`/`, and `^` in favor of ones that construct [`GenericNonlinearExpr`](@ref).
+
+This macro will error if the resulting expression does not produce a
+[`GenericNonlinearExpr`](@ref) because, for example, it is used on an expression
+that does not use the basic arithmetic operators.
+
+## When to use this macro
+
+In most cases, you should not use this macro.
+
+Use this macro only if the intended output type is a [`GenericNonlinearExpr`](@ref)
+and the regular macro calls destroy problem structure, or in rare cases, if the
+regular macro calls introduce a large amount of intermediate variables, for
+example, because they promote types to a common quadratic expression.
+
+## Example
+
+### Use-case one: preserve problem structure.
+
+```jldoctest
+julia> model = Model();
+
+julia> @variable(model, x);
+
+julia> @expression(model, (x - 0.1)^2)
+x² - 0.2 x + 0.010000000000000002
+
+julia> @expression(model, @force_nonlinear((x - 0.1)^2))
+(x - 0.1) ^ 2
+
+julia> (x - 0.1)^2
+x² - 0.2 x + 0.010000000000000002
+
+julia> @force_nonlinear((x - 0.1)^2)
+(x - 0.1) ^ 2
+```
+
+### Use-case two: reduce allocations
+
+In this example, we know that `x * 2.0 * (1 + x) * x` is going to construct a
+nonlinear expression.
+
+However, the default parsing first constructs:
+
+ * the [`GenericAffExpr`](@ref) `a = x * 2.0`,
+ * another [`GenericAffExpr`](@ref) `b = 1 + x`
+ * the [`GenericQuadExpr`](@ref) `c = a * b`
+ * a [`GenericNonlinearExpr`](@ref) `*(c, x)`
+
+In contrast, the modified parsing constructs:
+
+ * the [`GenericNonlinearExpr`](@ref) `a = GenericNonlinearExpr(:+, 1, x)`
+ * the [`GenericNonlinearExpr`](@ref) `GenericNonlinearExpr(:*, x, 2.0, a, x)`
+
+This results in significantly fewer allocations.
+
+```jldoctest
+julia> model = Model();
+
+julia> @variable(model, x);
+
+julia> @expression(model, x * 2.0 * (1 + x) * x)
+(2 x² + 2 x) * x
+
+julia> @expression(model, @force_nonlinear(x * 2.0 * (1 + x) * x))
+x * 2.0 * (1 + x) * x
+
+julia> @allocated @expression(model, x * 2.0 * (1 + x) * x)
+3200
+
+julia> @allocated @expression(model, @force_nonlinear(x * 2.0 * (1 + x) * x))
+640
+```
+"""
+macro force_nonlinear(expr)
+    error_fn = Containers.build_error_fn(:force_nonlinear, (expr,), __source__)
+    ret = MacroTools.postwalk(expr) do x
+        if Meta.isexpr(x, :call)
+            if x.args[1] == :+
+                return Expr(:call, _op_add, x.args[2:end]...)
+            elseif x.args[1] == :-
+                return Expr(:call, _op_sub, x.args[2:end]...)
+            elseif x.args[1] == :*
+                return Expr(:call, _op_mul, x.args[2:end]...)
+            elseif x.args[1] == :/
+                return Expr(:call, _op_div, x.args[2:end]...)
+            elseif x.args[1] == :^
+                return Expr(:call, _op_pow, x.args[2:end]...)
+            end
+        end
+        return x
+    end
+    return Expr(:call, _force_nonlinear, error_fn, esc(ret))
+end
+
+_force_nonlinear(::F, ret::GenericNonlinearExpr) where {F} = ret
+
+function _force_nonlinear(error_fn::F, ret::Any) where {F}
+    return error_fn(
+        "expression did not produce a `GenericNonlinearExpr`. Got a " *
+        "`$(typeof(ret))`: $(ret)",
+    )
+end

test/test_macros.jl

@@ -2358,4 +2358,26 @@ function test_op_or_short_circuit()
     return
 end
 
+function test_force_nonlinear()
+    model = Model()
+    @variable(model, x)
+    @test 1 + x isa AffExpr
+    @test @force_nonlinear(1 + x) isa GenericNonlinearExpr
+    @test 1 - x isa AffExpr
+    @test @force_nonlinear(1 - x) isa GenericNonlinearExpr
+    @test 2 * x isa AffExpr
+    @test @force_nonlinear(2 * x) isa GenericNonlinearExpr
+    @test x / 3 isa AffExpr
+    @test @force_nonlinear(x / 3) isa GenericNonlinearExpr
+    @test x^2 isa QuadExpr
+    @test @force_nonlinear(x^2) isa GenericNonlinearExpr
+    @test_throws_runtime(
+        ErrorException(
+            "In `@force_nonlinear(x)`: expression did not produce a `GenericNonlinearExpr`. Got a `$(typeof(x))`: $x",
+        ),
+        @force_nonlinear(x),
+    )
+    return
+end
+
 end  # module