diff --git a/docs/src/manual/nonlinear.md b/docs/src/manual/nonlinear.md index c097ebeb456..afd11e58876 100644 --- a/docs/src/manual/nonlinear.md +++ b/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 diff --git a/src/macros.jl b/src/macros.jl index f4a025778d8..041e49e5eec 100644 --- a/src/macros.jl +++ b/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 diff --git a/src/macros/@force_nonlinear.jl b/src/macros/@force_nonlinear.jl new file mode 100644 index 00000000000..5ea1dfde2a7 --- /dev/null +++ b/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 diff --git a/test/test_macros.jl b/test/test_macros.jl index 03652a0f679..9a144f7126d 100644 --- a/test/test_macros.jl +++ b/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