Expand lazy sums with explicit nonlinear self-loop bridges

README.md

@@ -41,6 +41,35 @@ To create a grouped sum, use `lazy_sum` as follows:
     lazy_sum(0.5 * R[j] * (u[i, j]^2) for i in 1:N, j in 1:p),
 )
 ```
+## Solve polynomial models with HiGHS
+
+Register the generator and scalar polynomial bridges to use affine and
+quadratic lazy sums with HiGHS:
+
+```julia
+using JuMP, GenOpt
+import HiGHS
+
+model = Model(HiGHS.Optimizer)
+add_bridge(model, GenOpt.FunctionGeneratorBridge)
+add_bridge(model, GenOpt.PolynomialFunctionBridge)
+@variable(model, x[1:4])
+@constraint(model, [i in 1:4], x[i] >= 0, container = ParametrizedArray)
+@constraint(model, lazy_sum(x[i] + 1 for i in 1:4) <= 12)
+@objective(model, Min, lazy_sum((x[i] - i)^2 for i in 1:4))
+optimize!(model)
+```
+
+GenOpt tracks known polynomial expressions using `PolynomialFunction{F}`, where
+`F` is the affine or quadratic output type. The shared builders accumulate lazy
+sum terms directly into that type when the optimizer needs it. HiGHS supports
+quadratic objectives, but not quadratic constraints. Ordinary, untagged
+`MOI.ScalarNonlinearFunction` inputs are not handled by these polynomial bridges.
+
+Polynomial lowering currently requires filters to refer only to the iterators
+bound by their sum. A filter depending on an enclosing generator raises an
+explicit error.
+
 ## Solve with ExaModels
 
 Install ExaModels and a solver such as MadNLP:

examples/opf/runtests.jl

@@ -37,8 +37,10 @@ const NONLINEAR = FunctionGenerator{MOI.ScalarNonlinearFunction}
     end
     # The reference bus angle, over the single reference bus.
     @test dims[(AFFINE, MOI.Zeros)] == [1]
-    # The 4 branch flow definitions and the 2 bus power balances, over 3 components each.
-    @test dims[(NONLINEAR, MOI.Zeros)] == fill(3, 6)
+    # The 4 branch flow definitions contain trigonometric terms.
+    @test dims[(NONLINEAR, MOI.Zeros)] == fill(3, 4)
+    # The 2 bus power balances contain affine sums and squared voltages.
+    @test dims[(QUADRATIC, MOI.Zeros)] == [3, 3]
     # `angmin[i] <= va[..] - va[..] <= angmax[i]` stays one interval constraint carrying
     # the per-branch bounds, rather than two inequalities per branch.
     @test dims[(AFFINE, MOI.HyperRectangle{Float64})] == [3]
@@ -53,9 +55,12 @@ end
     b = JuMP.backend(model)
     F = MOI.get(b, MOI.ObjectiveFunctionType())
     # The `lazy_sum` over the generators is not expanded into 3 cost terms.
-    @test F == SumGenerator{MOI.ScalarQuadraticFunction{Float64}}
+    @test F == GenOpt.PolynomialFunction{MOI.ScalarQuadraticFunction{Float64}}
     obj = MOI.get(b, MOI.ObjectiveFunction{F}())
-    @test prod(it -> length(it.values), obj.iterators) == 3
+    @test obj.func.head == :+
+    generator = only(obj.func.args)
+    @test generator isa SumGenerator{MOI.ScalarQuadraticFunction{Float64}}
+    @test prod(it -> length(it.values), generator.iterators) == 3
 end
 
 end  # module TestOPF

examples/quadrotor/runtests.jl

@@ -77,12 +77,26 @@ end
 @testset "objective keeps its generators" begin
     b = JuMP.backend(model)
     F = MOI.get(b, MOI.ObjectiveFunctionType())
-    @test F == MOI.ScalarNonlinearFunction
-    obj = MOI.get(b, MOI.ObjectiveFunction{F}())
+    @test F == PolynomialFunction{MOI.ScalarQuadraticFunction{Float64}}
+    obj = MOI.get(b, MOI.ObjectiveFunction{F}()).func
     # The three `lazy_sum` terms are summed without being expanded.
     @test obj.head == :+
-    @test all(a -> a isa SumGenerator, obj.args)
-    terms = [prod(it -> length(it.values), a.iterators) for a in obj.args]
+    # Additions may be nested, with polynomial degree tags around each sum.
+    pending = Any[obj]
+    sums = SumGenerator{MOI.ScalarQuadraticFunction{Float64}}[]
+    while !isempty(pending)
+        arg = pop!(pending)
+        if arg isa PolynomialFunction
+            push!(pending, arg.func)
+        elseif arg isa MOI.ScalarNonlinearFunction
+            @test arg.head == :+
+            append!(pending, arg.args)
+        else
+            @test arg isa SumGenerator{MOI.ScalarQuadraticFunction{Float64}}
+            push!(sums, arg)
+        end
+    end
+    terms = [prod(it -> length(it.values), a.iterators) for a in sums]
     # `N * p` control terms, `N * n` stage terms and `n` terminal terms.
     # Older JuMP versions may reorder terms when flattening the objective sum.
     @test sort(terms) == sort([N * 4, N * n, n])

ext/GenOptExaModelsExt.jl

@@ -7,7 +7,11 @@ module GenOptExaModelsExt
 
 import ExaModels
 import GenOpt:
-    FunctionGenerator, SumGenerator, ContiguousArrayOfVariables, IteratorIndex
+    FunctionGenerator,
+    SumGenerator,
+    PolynomialFunction,
+    ContiguousArrayOfVariables,
+    IteratorIndex
 import MathOptInterface as MOI
 
 # The triggers of this extension are a strict superset of the ones of
@@ -18,6 +22,34 @@ const ExaModelsMOI = Base.get_extension(ExaModels, :ExaModelsMOI)
 
 # Objective
 
+function MOI.supports(
+    ::ExaModelsMOI.Optimizer,
+    ::MOI.ObjectiveFunction{<:PolynomialFunction},
+)
+    return true
+end
+
+function MOI.set(
+    model::ExaModelsMOI.Optimizer,
+    ::MOI.ObjectiveFunction{F},
+    f::F,
+) where {F<:PolynomialFunction}
+    # Unwrap the degree tag, keeping generator templates and iterator data.
+    return MOI.set(
+        model,
+        MOI.ObjectiveFunction{MOI.ScalarNonlinearFunction}(),
+        f.func,
+    )
+end
+
+function ExaModelsMOI.update_bin!(
+    bins::Vector{ExaModelsMOI.Bin},
+    fn::ExaModelsMOI.AbstractBin,
+    f::PolynomialFunction,
+)
+    return ExaModelsMOI.update_bin!(bins, fn, f.func)
+end
+
 function MOI.supports(
     ::ExaModelsMOI.Optimizer,
     ::MOI.ObjectiveFunction{<:SumGenerator},
@@ -103,6 +135,11 @@ function _data(iterators)
 end
 
 exagen(α::Number, _) = α
+exagen(f::PolynomialFunction, offsets) = exagen(f.func, offsets)
+
+function ExaModelsMOI._exafy(f::PolynomialFunction, data::Tuple, var_to_data)
+    return ExaModelsMOI._exafy(f.func, data, var_to_data)
+end
 
 function exagen(f::MOI.ScalarNonlinearFunction, offsets)
     if f.head != :getindex

src/Bridges/Bridges.jl

@@ -4,3 +4,4 @@
 # in the LICENSE.md file or at https://opensource.org/licenses/MIT.
 
 include("Constraint/Constraint.jl")
+include("polynomial_function.jl")

src/Bridges/Constraint/Constraint.jl

@@ -4,3 +4,4 @@
 # in the LICENSE.md file or at https://opensource.org/licenses/MIT.
 
 include("function_generator.jl")
+include("polynomial_function.jl")