diff --git a/Project.toml b/Project.toml index 29165305..4fa84201 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" [weakdeps] +GenOpt = "f2c049d8-7489-4223-990c-4f1c121a4cde" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -19,6 +20,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] +ExaModelsGenOpt = ["GenOpt", "MathOptInterface"] ExaModelsJuMP = "JuMP" ExaModelsKernelAbstractions = "KernelAbstractions" ExaModelsMOI = "MathOptInterface" @@ -29,6 +31,7 @@ ExaModelsSpecialFunctions = "SpecialFunctions" [compat] Adapt = "4" +GenOpt = "0.2.2" JuMP = "1" KernelAbstractions = "0.9" MathOptInterface = "1.19" diff --git a/ext/ExaModelsGenOpt.jl b/ext/ExaModelsGenOpt.jl new file mode 100644 index 00000000..a0f86845 --- /dev/null +++ b/ext/ExaModelsGenOpt.jl @@ -0,0 +1,127 @@ +module ExaModelsGenOpt + +import ExaModels +import GenOpt +import GenOpt: FunctionGenerator, SumGenerator, ContiguousArrayOfVariables, IteratorIndex, Iterator +import MathOptInterface as MOI + +# Mark GenOpt function types as extension types +ExaModels.is_extension_type(::Type{<:FunctionGenerator}) = true +ExaModels.is_extension_type(::Type{<:SumGenerator}) = true + +function _map_indices(index_map, f::MOI.ScalarNonlinearFunction) + args = Any[_map_indices(index_map, arg) for arg in f.args] + return MOI.ScalarNonlinearFunction(f.head, args) +end + +function _map_indices(index_map, array::ContiguousArrayOfVariables) + first_src = MOI.VariableIndex(array.offset + 1) + first_dest = _map_variable(index_map, first_src) + return ContiguousArrayOfVariables(first_dest.value - 1, array.size) +end + +_map_indices(::Any, arg) = arg + +_map_variable(index_map::MOI.Utilities.IndexMap, variable) = index_map[variable] +_map_variable(index_map::Function, variable) = index_map(variable) + +function MOI.Utilities.map_indices( + index_map::MOI.Utilities.IndexMap, + f::FunctionGenerator{F}, +) where {F} + return FunctionGenerator{F}(_map_indices(index_map, f.func), f.iterators) +end + + +function MOI.Utilities.map_indices( + index_map::Function, + f::FunctionGenerator{F}, +) where {F} + return FunctionGenerator{F}(_map_indices(index_map, f.func), f.iterators) +end + +function MOI.Utilities.map_indices( + index_map::Function, + f::SumGenerator{F}, +) where {F} + return SumGenerator{F}(_map_indices(index_map, f.func), f.iterators) +end + +function MOI.Utilities.map_indices( + index_map::MOI.Utilities.IndexMap, + f::SumGenerator{F}, +) where {F} + return SumGenerator{F}(_map_indices(index_map, f.func), f.iterators) +end + +# Handle SumGenerator in objective expressions +function ExaModels.exafy_extension_obj_arg(m::SumGenerator) + return _exagen(m.func, m.iterators) +end + +function ExaModels.add_extra_constraint!(model, f::FunctionGenerator, s) + exa_moi = Base.get_extension(ExaModels, :ExaModelsMOI) + row = length(model.lcon) + 1 + expr, pars = _exagen(f.func, f.iterators) + indexed = ExaModels.DataIndexed(ExaModels.DataSource(), length(first(pars)) + 1) + data = [(p..., row + i - 1) for (i, p) in enumerate(pars)] + push!(model.cons, exa_moi.Bin(indexed => expr, data)) + append!(model.lcon, _lower_bounds(s, eltype(model.lcon))) + append!(model.ucon, _upper_bounds(s, eltype(model.ucon))) + return MOI.ConstraintIndex{typeof(f),typeof(s)}(row) +end + +# Convert GenOpt expression trees to ExaModels format + +exagen(α::Number, _) = α + +function exagen(f::MOI.ScalarNonlinearFunction, offsets) + if f.head == :getindex + v = f.args[1] + if v isa ContiguousArrayOfVariables + idx = exagen(f.args[2], offsets) + if !iszero(v.offset) + idx = v.offset + idx + end + cp = cumprod(v.size) + for i in 3:length(f.args) + idx += cp[i - 2] * (exagen(f.args[i], offsets) - 1) + end + return ExaModels.Var(idx) + elseif v isa IteratorIndex + @assert length(f.args) == 2 + @assert f.args[2] isa Integer + if isnothing(offsets) + @assert isone(f.args[2]) + return ExaModels.DataSource() + else + return ExaModels.DataIndexed(ExaModels.DataSource(), offsets[v.value] + f.args[2]) + end + else + error("Unexpected the first operand of `getindex` to be of type `$(typeof(v))`") + end + else + op = getfield(MOI.Nonlinear, f.head) + return op((exagen(e, offsets) for e in f.args)...) + end +end + +function _exagen(func::MOI.ScalarNonlinearFunction, iterators) + lengths = map(it -> length(first(it.values)), iterators) + cs = [0; cumsum(lengths)[1:(end - 1)]] + pars = vec( + map(Base.Iterators.ProductIterator(ntuple(i -> iterators[i].values, length(iterators)))) do I + reduce((i, j) -> tuple(i..., j...), I) + end + ) + expr = exagen(func, cs) + return expr, pars +end + +# Bound helpers for vector sets used by FunctionGenerator constraints +_lower_bounds(s::Union{MOI.Zeros,MOI.Nonnegatives}, T) = fill(zero(T), MOI.dimension(s)) +_lower_bounds(s::MOI.Nonpositives, T) = fill(typemin(T), MOI.dimension(s)) +_upper_bounds(s::Union{MOI.Zeros,MOI.Nonpositives}, T) = fill(zero(T), MOI.dimension(s)) +_upper_bounds(s::MOI.Nonnegatives, T) = fill(typemax(T), MOI.dimension(s)) + +end # module diff --git a/ext/ExaModelsMOI.jl b/ext/ExaModelsMOI.jl index b51418f7..1a2d2b39 100644 --- a/ext/ExaModelsMOI.jl +++ b/ext/ExaModelsMOI.jl @@ -202,6 +202,16 @@ function update_bin!( return bins end +function update_bin!(bins::Vector{Bin}, fn::AbstractBin, f) + if !ExaModels.is_extension_type(typeof(f)) + throw(MOI.UnsupportedAttribute(MOI.ObjectiveFunction{typeof(f)}())) + end + @assert fn isa ObjectiveBin + head, data = ExaModels.exafy_extension_obj_arg(f) + push!(bins, Bin(head, data)) + return bins +end + # _exafy # This method is used for objective constants. @@ -482,6 +492,13 @@ function MOI.supports( return true end +function MOI.supports( + ::Optimizer, + ::MOI.ObjectiveFunction{F}, +) where {F} + return ExaModels.is_extension_type(F) +end + function MOI.set( model::Optimizer{T}, ::MOI.ObjectiveFunction{F}, @@ -500,6 +517,16 @@ function MOI.set( return end +function MOI.set(model::Optimizer, ::MOI.ObjectiveFunction{F}, f::F) where {F} + if !ExaModels.is_extension_type(F) + throw(MOI.UnsupportedAttribute(MOI.ObjectiveFunction{F}())) + end + empty!(model.objs) + head, data = ExaModels.exafy_extension_obj_arg(f) + push!(model.objs, Bin(head, data)) + return +end + # MOI.add_variable function MOI.add_variable(model::Optimizer{T}) where {T} @@ -546,6 +573,14 @@ function MOI.supports_constraint( return true end +function MOI.supports_constraint( + ::Optimizer, + ::Type{F}, + ::Type{S}, +) where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet} + return ExaModels.is_extension_type(F) +end + function _update_bound(model::Optimizer, col::Int, set::MOI.GreaterThan) model.lvar[col] = set.lower return @@ -655,6 +690,17 @@ function MOI.add_constraint( return MOI.ConstraintIndex{typeof(f),typeof(s)}(row) end +function MOI.add_constraint( + model::Optimizer, + f::F, + s::S, +) where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet} + if !ExaModels.is_extension_type(F) + throw(MOI.UnsupportedConstraint{F,S}()) + end + return ExaModels.add_extra_constraint!(model, f, s) +end + function to_exacore(model::Optimizer{T}, backend) where {T} c = ExaModels.ExaCore( T; diff --git a/src/ExaModels.jl b/src/ExaModels.jl index e3b9b412..c4f8b3f1 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -61,6 +61,7 @@ include("oracle.jl") include("utils.jl") include("tags.jl") include("two_stage.jl") +include("wrapper.jl") export ExaModel, ExaCore, diff --git a/src/wrapper.jl b/src/wrapper.jl new file mode 100644 index 00000000..b308b362 --- /dev/null +++ b/src/wrapper.jl @@ -0,0 +1,42 @@ +# Extension points used by ExaModelsMOI and ExaModelsGenOpt extensions + +""" + copy_extra_constraints!(c, moim, var_to_idx, con_to_idx, T) + +Hook for extensions to add extra constraint types after standard MOI constraints +are processed. Default is a no-op, defined in ExaModelsMOI. +""" +function copy_extra_constraints! end + +""" + is_extension_type(::Type{F}) -> Bool + +Return `true` if `F` is a function type handled by an extension. +Used by `check_supported` and `supports_constraint` to whitelist extension types. +""" +function is_extension_type end +is_extension_type(::Type) = false + +""" + exafy_extension_obj_arg(m, var_to_idx) -> Union{Nothing, Tuple} + +Try to convert an objective function argument `m` to an `(expr, pars)` tuple +for ExaModels. `var_to_idx` maps `MOI.VariableIndex` to `(type, idx)` named tuples. +Returns `nothing` if the type is not handled by any extension. +""" +function exafy_extension_obj_arg end + +""" + add_extra_constraint!(model, f, s) + +Add a constraint whose function type is implemented by an extension. +""" +function add_extra_constraint! end + +""" + op(s::Symbol) + +Map a Symbol to the corresponding Julia function. Used by both ExaModelsMOI +and ExaModelsGenOpt for expression tree conversion. +""" +function op end diff --git a/test/GenOptTest/GenOptTest.jl b/test/GenOptTest/GenOptTest.jl new file mode 100644 index 00000000..c707e827 --- /dev/null +++ b/test/GenOptTest/GenOptTest.jl @@ -0,0 +1,153 @@ +module GenOptTest + +using Test, JuMP, ExaModels, MadNLP, GenOpt + +function quadrotor_test() + container = GenOpt.ParametrizedArray + + N = 3 + n = 9 + p = 4 + dt = 0.1 + + x0_val = zeros(n) + xf = [1.0, 0, 0, 1, 0, 0, 0, 0, 0] + + Q = zeros(n) + Q[1] = 1.0; Q[4] = 1.0 + Qf = ones(n) + R = ones(p) + + g = 9.81 + + itr1 = [(i, j, xf[j]) for i in 1:N for j in 1:n if Q[j] != 0] + itr2 = [(j, xf[j]) for j in 1:n] + + model = Model() + @variable(model, x[1:(N + 1), 1:n]) + @variable(model, u[1:N, 1:p]) + + @constraint(model, start[i in 1:n], x[1, i] == x0_val[i], container = container) + + @constraint(model, [i in 1:N], x[i + 1, 1] == x[i, 1] + x[i, 2] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 2] == x[i, 2] + (u[i, 1]) * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 3] == x[i, 3] + x[i, 4] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 4] == x[i, 4] + (u[i, 2]) * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 5] == x[i, 5] + x[i, 6] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 6] == x[i, 6] + (u[i, 3] - g) * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 7] == x[i, 7] + u[i, 1] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 8] == x[i, 8] + u[i, 2] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 9] == x[i, 9] + u[i, 4] * dt, container = container) + + @objective( + model, Min, + lazy_sum(0.5 * R[j] * (u[i, j]^2) for i in 1:N, j in 1:p) + + lazy_sum(0.5 * Q[it[2]] * (x[it[1], it[2]] - it[3])^2 for it in itr1) + + lazy_sum(0.5 * Qf[it[1]] * (x[N + 1, it[1]] - it[2])^2 for it in itr2), + ) + + return model +end + +function _get_optimizer(model) + be = backend(model) + return be.optimizer.model +end + +function parametric_genopt_test() + # Same structure as quadrotor_test but with MOI parameters declared first, + # which shifts MOI variable indices relative to ExaModels internal indices. + # This tests that the GenOpt extension does NOT assume var_to_idx is the identity. + container = GenOpt.ParametrizedArray + + N = 3 + n = 9 + p = 4 + dt = 0.1 + + x0_val = zeros(n) + xf = [1.0, 0, 0, 1, 0, 0, 0, 0, 0] + + Q = zeros(n) + Q[1] = 1.0; Q[4] = 1.0 + Qf = ones(n) + R = ones(p) + + g = 9.81 + + itr1 = [(i, j, xf[j]) for i in 1:N for j in 1:n if Q[j] != 0] + itr2 = [(j, xf[j]) for j in 1:n] + + model = Model() + + # Parameters come first, shifting MOI variable indices by 2 + @variable(model, par1 in MOI.Parameter(2.0)) + @variable(model, par2 in MOI.Parameter(3.0)) + + @variable(model, x[1:(N + 1), 1:n]) + @variable(model, u[1:N, 1:p]) + + @constraint(model, start[i in 1:n], x[1, i] == x0_val[i], container = container) + + @constraint(model, [i in 1:N], x[i + 1, 1] == x[i, 1] + x[i, 2] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 2] == x[i, 2] + (u[i, 1]) * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 3] == x[i, 3] + x[i, 4] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 4] == x[i, 4] + (u[i, 2]) * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 5] == x[i, 5] + x[i, 6] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 6] == x[i, 6] + (u[i, 3] - g) * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 7] == x[i, 7] + u[i, 1] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 8] == x[i, 8] + u[i, 2] * dt, container = container) + @constraint(model, [i in 1:N], x[i + 1, 9] == x[i, 9] + u[i, 4] * dt, container = container) + + @objective( + model, Min, + lazy_sum(0.5 * R[j] * (u[i, j]^2) for i in 1:N, j in 1:p) + + lazy_sum(0.5 * Q[it[2]] * (x[it[1], it[2]] - it[3])^2 for it in itr1) + + lazy_sum(0.5 * Qf[it[1]] * (x[N + 1, it[1]] - it[2])^2 for it in itr2), + ) + + return model +end + +function runtests() + return @testset "GenOpt extension test" begin + @testset "Quadrotor with ExaModels.Optimizer" begin + model = quadrotor_test() + + set_optimizer(model, () -> ExaModels.Optimizer(MadNLP.madnlp)) + set_optimizer_attribute(model, "print_level", MadNLP.ERROR) + optimize!(model) + + obj = objective_value(model) + @test isapprox(obj, 8.1797, atol = 1.0e-3) + + # Verify that FunctionGenerator structure is retained in ExaModels + # (not dismantled into individual scalar constraints). + # The test has 10 @constraint calls with `container = GenOpt.ParametrizedArray`: + # - 1 start constraint with n=9 iterations + # - 9 dynamics constraints each with N=3 iterations + # Plus 1 empty block for standard (non-generator) constraints. + # If structure were lost, we'd see 36 entries each with 1 iteration. + optimizer = _get_optimizer(model) + @test length(optimizer.lcon) == 36 + @test length(optimizer.cons) == 10 + itr_lengths = sort([length(c.data) for c in optimizer.cons]) + @test itr_lengths == [3, 3, 3, 3, 3, 3, 3, 3, 3, 9] + end + + @testset "GenOpt with MOI parameters" begin + # Tests that var_to_idx is not assumed to be the identity. + # Parameters before variables shift MOI indices. + model = parametric_genopt_test() + + set_optimizer(model, () -> ExaModels.Optimizer(MadNLP.madnlp)) + set_optimizer_attribute(model, "print_level", MadNLP.ERROR) + optimize!(model) + + # Same problem as quadrotor, should give same objective + @test isapprox(objective_value(model), 8.1797, atol = 1.0e-3) + end + end +end + +end # module diff --git a/test/Project.toml b/test/Project.toml index 043fb9ec..066c0897 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -2,6 +2,7 @@ ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2" ExaPowerIO = "14903efe-9500-4d7f-a589-7ab7e15da6de" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +GenOpt = "f2c049d8-7489-4223-990c-4f1c121a4cde" Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" JuliaC = "acedd4c2-ced6-4a15-accc-2607eb759ba2" diff --git a/test/runtests.jl b/test/runtests.jl index cc22b2da..3511dda5 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -27,6 +27,7 @@ include("PrettyPrintTest.jl") include("ConcreteModeTest.jl") # include("OptimalControlTest/OptimalControlTest.jl") include("OracleTest/OracleTest.jl") +include("GenOptTest/GenOptTest.jl") @testset verbose = true "ExaModels test" begin @info "Running Argument Test" @@ -64,6 +65,9 @@ include("OracleTest/OracleTest.jl") @info "Running Oracle Test" OracleTest.runtests() + + @info "Running GenOpt Test" + GenOptTest.runtests() end # Force full GC before Julia exits so that OpenCL/PoCL objects are finalized