From 6087e3bd01bdb1895300aa1e8a558a8bbac49cd1 Mon Sep 17 00:00:00 2001 From: odow Date: Tue, 16 Apr 2024 12:26:37 +1200 Subject: [PATCH 1/8] Add @nl macro for modifying how expressions are parsed --- src/macros.jl | 3 +- src/macros/{@NL.jl => @NLlegacy.jl} | 0 src/macros/@nl_macro.jl | 98 +++++++++++++++++++++++++++++ 3 files changed, 100 insertions(+), 1 deletion(-) rename src/macros/{@NL.jl => @NLlegacy.jl} (100%) create mode 100644 src/macros/@nl_macro.jl diff --git a/src/macros.jl b/src/macros.jl index f4a025778d8..a28f36f8c27 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -471,4 +471,5 @@ include("macros/@objective.jl") include("macros/@expression.jl") include("macros/@constraint.jl") include("macros/@variable.jl") -include("macros/@NL.jl") +include("macros/@nl_macro.jl") +include("macros/@NLlegacy.jl") diff --git a/src/macros/@NL.jl b/src/macros/@NLlegacy.jl similarity index 100% rename from src/macros/@NL.jl rename to src/macros/@NLlegacy.jl diff --git a/src/macros/@nl_macro.jl b/src/macros/@nl_macro.jl new file mode 100644 index 00000000000..3977580abad --- /dev/null +++ b/src/macros/@nl_macro.jl @@ -0,0 +1,98 @@ +# 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(/, :/) + +""" + @nl(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). + +## 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, @nl((x - 0.1)^2)) +(x - 0.1) ^ 2.0 + +julia> (x - 0.1)^2 +x² - 0.2 x + 0.010000000000000002 + +julia> @nl((x - 0.1)^2) +(x - 0.1) ^ 2.0 +``` + +### 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> @allocated @expression(model, x * 2.0 * (1 + x) * x) +3200 + +julia> @allocated @expression(model, @nl(x * 2.0 * (1 + x) * x)) +640 +``` +""" +macro nl(expr) + 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]...) + end + end + return x + end + return esc(ret) +end From dba20fc5dc613a6210974a2651a4564d5e500e88 Mon Sep 17 00:00:00 2001 From: Oscar Dowson Date: Tue, 16 Apr 2024 12:54:38 +1200 Subject: [PATCH 2/8] Update @nl_macro.jl --- src/macros/@nl_macro.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/macros/@nl_macro.jl b/src/macros/@nl_macro.jl index 3977580abad..cc0b0c2f793 100644 --- a/src/macros/@nl_macro.jl +++ b/src/macros/@nl_macro.jl @@ -72,6 +72,12 @@ julia> model = Model(); julia> @variable(model, x); +julia> @expression(model, x * 2.0 * (1 + x) * x) +(2 x² + 2 x) * x + +julia> @expression(model, @nl(x * 2.0 * (1 + x) * x)) +x * 2.0 * (1 + x) * x + julia> @allocated @expression(model, x * 2.0 * (1 + x) * x) 3200 From 900eeff9cdfc00f0694e981c6cffd5559bac7d91 Mon Sep 17 00:00:00 2001 From: odow Date: Thu, 18 Apr 2024 10:21:54 +1200 Subject: [PATCH 3/8] Update --- src/macros.jl | 9 +++------ src/macros/{@NLlegacy.jl => @NL.jl} | 0 src/macros/{@nl_macro.jl => @nonlinear.jl} | 12 ++++++------ 3 files changed, 9 insertions(+), 12 deletions(-) rename src/macros/{@NLlegacy.jl => @NL.jl} (100%) rename src/macros/{@nl_macro.jl => @nonlinear.jl} (91%) diff --git a/src/macros.jl b/src/macros.jl index a28f36f8c27..71aa93cb48f 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -467,9 +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_macro.jl") -include("macros/@NLlegacy.jl") +for file in readdir("macros") + include(joinpath("macros", file)) +end diff --git a/src/macros/@NLlegacy.jl b/src/macros/@NL.jl similarity index 100% rename from src/macros/@NLlegacy.jl rename to src/macros/@NL.jl diff --git a/src/macros/@nl_macro.jl b/src/macros/@nonlinear.jl similarity index 91% rename from src/macros/@nl_macro.jl rename to src/macros/@nonlinear.jl index cc0b0c2f793..1be265aa67c 100644 --- a/src/macros/@nl_macro.jl +++ b/src/macros/@nonlinear.jl @@ -9,7 +9,7 @@ const _op_mul = NonlinearOperator(*, :*) const _op_div = NonlinearOperator(/, :/) """ - @nl(expr) + @nonlinear(expr) Change the parsing of `expr` to construct [`GenericNonlinearExpr`](@ref) instead of [`GenericAffExpr`](@ref) or [`GenericQuadExpr`](@ref). @@ -38,13 +38,13 @@ julia> @variable(model, x); julia> @expression(model, (x - 0.1)^2) x² - 0.2 x + 0.010000000000000002 -julia> @expression(model, @nl((x - 0.1)^2)) +julia> @expression(model, @nonlinear((x - 0.1)^2)) (x - 0.1) ^ 2.0 julia> (x - 0.1)^2 x² - 0.2 x + 0.010000000000000002 -julia> @nl((x - 0.1)^2) +julia> @nonlinear((x - 0.1)^2) (x - 0.1) ^ 2.0 ``` @@ -75,17 +75,17 @@ julia> @variable(model, x); julia> @expression(model, x * 2.0 * (1 + x) * x) (2 x² + 2 x) * x -julia> @expression(model, @nl(x * 2.0 * (1 + x) * x)) +julia> @expression(model, @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, @nl(x * 2.0 * (1 + x) * x)) +julia> @allocated @expression(model, @nonlinear(x * 2.0 * (1 + x) * x)) 640 ``` """ -macro nl(expr) +macro nonlinear(expr) ret = MacroTools.postwalk(expr) do x if Meta.isexpr(x, :call) if x.args[1] == :+ From b49ee11dd883fdad84a3dd3bc03b25b5bc5d3f95 Mon Sep 17 00:00:00 2001 From: odow Date: Thu, 18 Apr 2024 10:35:12 +1200 Subject: [PATCH 4/8] Update --- src/macros.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index 71aa93cb48f..041e49e5eec 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -467,6 +467,6 @@ function _plural_macro_code(model, block, macro_sym) return code end -for file in readdir("macros") - include(joinpath("macros", file)) +for file in readdir(joinpath(@__DIR__, "macros")) + include(joinpath(@__DIR__, "macros", file)) end From e9f8eb67a22f0a11fc12c6a2f66164b6af904427 Mon Sep 17 00:00:00 2001 From: odow Date: Mon, 29 Apr 2024 12:38:29 +1200 Subject: [PATCH 5/8] Update --- .../{@nonlinear.jl => @force_nonlinear.jl} | 32 +++++++++++++------ test/test_macros.jl | 22 +++++++++++++ 2 files changed, 45 insertions(+), 9 deletions(-) rename src/macros/{@nonlinear.jl => @force_nonlinear.jl} (73%) 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 From 33dbb24bfb4cac76e3217baef24f3395f42c72d0 Mon Sep 17 00:00:00 2001 From: odow Date: Mon, 29 Apr 2024 12:59:21 +1200 Subject: [PATCH 6/8] Update --- src/macros/@force_nonlinear.jl | 4 ++-- test/test_macros.jl | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/macros/@force_nonlinear.jl b/src/macros/@force_nonlinear.jl index 1b3be60082a..530bccf7e80 100644 --- a/src/macros/@force_nonlinear.jl +++ b/src/macros/@force_nonlinear.jl @@ -44,13 +44,13 @@ 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.0 +(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.0 +(x - 0.1) ^ 2 ``` ### Use-case two: reduce allocations diff --git a/test/test_macros.jl b/test/test_macros.jl index 24ca9d20747..0ea3fb23efc 100644 --- a/test/test_macros.jl +++ b/test/test_macros.jl @@ -2361,16 +2361,16 @@ end function test_force_nonlinear() model = Model() @variable(model, x) - @test (1 + x) isa AffExpr + @test 1 + x isa AffExpr @test @force_nonlinear(1 + x) isa GenericNonlinearExpr - @test (1 - x) isa AffExpr + @test 1 - x isa AffExpr @test @force_nonlinear(1 - x) isa GenericNonlinearExpr - @test (2 * x) isa AffExpr + @test 2 * x isa AffExpr @test @force_nonlinear(2 * x) isa GenericNonlinearExpr - @test (x / 3) isa AffExpr + @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 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", From 75415a3d3d1589698cd6369219dbb997b8de8028 Mon Sep 17 00:00:00 2001 From: odow Date: Mon, 29 Apr 2024 15:16:44 +1200 Subject: [PATCH 7/8] Update --- docs/src/manual/nonlinear.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/docs/src/manual/nonlinear.md b/docs/src/manual/nonlinear.md index c097ebeb456..dcba37b45d5 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 over-ride 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 From 9ab125e8793e105b4594ba1b2fb2342b7c1275ec Mon Sep 17 00:00:00 2001 From: odow Date: Fri, 3 May 2024 08:56:18 +1200 Subject: [PATCH 8/8] Update --- docs/src/manual/nonlinear.md | 2 +- src/macros/@force_nonlinear.jl | 17 ++++++++++------- test/test_macros.jl | 2 +- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/docs/src/manual/nonlinear.md b/docs/src/manual/nonlinear.md index dcba37b45d5..afd11e58876 100644 --- a/docs/src/manual/nonlinear.md +++ b/docs/src/manual/nonlinear.md @@ -370,7 +370,7 @@ x² - 0.2 x + 0.010000000000000002 julia> typeof(f) QuadExpr (alias for GenericQuadExpr{Float64, GenericVariableRef{Float64}}) ``` -To over-ride this behavior, use the [`@force_nonlinear`](@ref) macro: +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 diff --git a/src/macros/@force_nonlinear.jl b/src/macros/@force_nonlinear.jl index 530bccf7e80..5ea1dfde2a7 100644 --- a/src/macros/@force_nonlinear.jl +++ b/src/macros/@force_nonlinear.jl @@ -108,11 +108,14 @@ macro force_nonlinear(expr) end return x end - return quote - r = $(esc(ret)) - if !(r isa $GenericNonlinearExpr) - $error_fn("expression did not produce a GenericNonlinearExpr") - end - r - 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 0ea3fb23efc..9a144f7126d 100644 --- a/test/test_macros.jl +++ b/test/test_macros.jl @@ -2373,7 +2373,7 @@ function test_force_nonlinear() @test @force_nonlinear(x^2) isa GenericNonlinearExpr @test_throws_runtime( ErrorException( - "In `@force_nonlinear(x)`: expression did not produce a GenericNonlinearExpr", + "In `@force_nonlinear(x)`: expression did not produce a `GenericNonlinearExpr`. Got a `$(typeof(x))`: $x", ), @force_nonlinear(x), )