diff --git a/src/macros/@nonlinear.jl b/src/macros/@force_nonlinear.jl similarity index 73% rename from src/macros/@nonlinear.jl rename to src/macros/@force_nonlinear.jl index 1be265aa67c..1b3be60082a 100644 --- a/src/macros/@nonlinear.jl +++ b/src/macros/@force_nonlinear.jl @@ -7,15 +7,20 @@ const _op_add = NonlinearOperator(+, :+) const _op_sub = NonlinearOperator(-, :-) const _op_mul = NonlinearOperator(*, :*) const _op_div = NonlinearOperator(/, :/) +const _op_pow = NonlinearOperator(^, :^) """ - @nonlinear(expr) + @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 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 @@ -38,13 +43,13 @@ julia> @variable(model, x); julia> @expression(model, (x - 0.1)^2) x² - 0.2 x + 0.010000000000000002 -julia> @expression(model, @nonlinear((x - 0.1)^2)) +julia> @expression(model, @force_nonlinear((x - 0.1)^2)) (x - 0.1) ^ 2.0 julia> (x - 0.1)^2 x² - 0.2 x + 0.010000000000000002 -julia> @nonlinear((x - 0.1)^2) +julia> @force_nonlinear((x - 0.1)^2) (x - 0.1) ^ 2.0 ``` @@ -75,17 +80,18 @@ julia> @variable(model, x); julia> @expression(model, x * 2.0 * (1 + x) * x) (2 x² + 2 x) * x -julia> @expression(model, @nonlinear(x * 2.0 * (1 + 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, @nonlinear(x * 2.0 * (1 + x) * x)) +julia> @allocated @expression(model, @force_nonlinear(x * 2.0 * (1 + x) * x)) 640 ``` """ -macro nonlinear(expr) +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] == :+ @@ -96,9 +102,17 @@ macro nonlinear(expr) 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 esc(ret) + return quote + r = $(esc(ret)) + if !(r isa $GenericNonlinearExpr) + $error_fn("expression did not produce a GenericNonlinearExpr") + end + r + end end diff --git a/test/test_macros.jl b/test/test_macros.jl index 03652a0f679..24ca9d20747 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", + ), + @force_nonlinear(x), + ) + return +end + end # module