diff --git a/Project.toml b/Project.toml index 7ed1bb39..c5ae051e 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" Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -23,6 +24,7 @@ oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" # oneAPI version >= 2.6 is neede # OptimalControl = "5f98b655-cc9a-415a-b60e-744165666948" [extensions] +ExaModelsGenOpt = ["GenOpt", "MathOptInterface"] ExaModelsIpopt = ["MathOptInterface", "NLPModelsIpopt"] ExaModelsJuMP = "JuMP" ExaModelsKernelAbstractions = "KernelAbstractions" @@ -36,6 +38,7 @@ ExaModelsSpecialFunctions = "SpecialFunctions" [compat] Adapt = "4" +GenOpt = "0.2" Ipopt = "1.11" JuMP = "1" KernelAbstractions = "0.9" @@ -48,4 +51,4 @@ OpenCL = "0.10" SolverCore = "0.3" SpecialFunctions = "2" julia = "1.9" -oneAPI = "2" \ No newline at end of file +oneAPI = "2" diff --git a/ext/ExaModelsGenOpt.jl b/ext/ExaModelsGenOpt.jl new file mode 100644 index 00000000..4f6dec04 --- /dev/null +++ b/ext/ExaModelsGenOpt.jl @@ -0,0 +1,99 @@ +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 + +# Handle SumGenerator in objective expressions +function ExaModels.exafy_extension_obj_arg(m::SumGenerator, var_to_idx) + return _exagen(m.func, m.iterators, var_to_idx) +end + +# Hook to process FunctionGenerator constraints after standard constraints +function ExaModels.copy_extra_constraints!(c, moim, var_to_idx, con_to_idx, T) + con_types = MOI.get(moim, MOI.ListOfConstraintTypesPresent()) + for (F, S) in con_types + F <: FunctionGenerator || continue + cis = MOI.get(moim, MOI.ListOfConstraintIndices{F, S}()) + c = _copy_generator_constraints!(c, moim, cis, var_to_idx, con_to_idx, T, S) + end + return c +end + +function _copy_generator_constraints!(c, moim, cis, var_to_idx, con_to_idx, T, ::Type{S}) where {S} + for ci in cis + func = MOI.get(moim, MOI.ConstraintFunction(), ci) + set = MOI.get(moim, MOI.ConstraintSet(), ci) + con_to_idx[ci] = c.ncon + expr, pars = _exagen(func.func, func.iterators, var_to_idx) + c, _ = ExaModels.add_con(c, expr for p in pars; lcon = _lower_bounds(set, T), ucon = _upper_bounds(set, T)) + end + return c +end + +# Convert GenOpt expression trees to ExaModels format + +exagen(α::Number, _, _) = α + +function exagen(f::MOI.ScalarNonlinearFunction, offsets, var_to_idx) + if f.head == :getindex + v = f.args[1] + if v isa ContiguousArrayOfVariables + idx = exagen(f.args[2], offsets, var_to_idx) + # Translate MOI-space offset to ExaModels-space offset using var_to_idx + first_moi_vi = MOI.VariableIndex(v.offset + 1) + exa_offset = var_to_idx[first_moi_vi].idx - 1 + if !iszero(exa_offset) + idx = exa_offset + idx + end + cp = cumprod(v.size) + for i in 3:length(f.args) + idx += cp[i - 2] * (exagen(f.args[i], offsets, var_to_idx) - 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 + return ExaModels.op(f.head)((exagen(e, offsets, var_to_idx) for e in f.args)...) + end +end + +function _exagen(func::MOI.ScalarNonlinearFunction, iterators, var_to_idx) + lengths = map(it -> length(first(it.values)), iterators) + if length(lengths) == 1 && lengths[] == 1 + cs = nothing + pars = only.(iterators[].values) + else + 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 + ) + end + expr = exagen(func, cs, var_to_idx) + return expr, pars +end + +# Bound helpers for vector sets used by FunctionGenerator constraints +_lower_bounds(::Union{MOI.Zeros, MOI.Nonnegatives}, T) = zero(T) +_lower_bounds(::MOI.Nonpositives, T) = typemin(T) +_upper_bounds(::Union{MOI.Zeros, MOI.Nonpositives}, T) = zero(T) +_upper_bounds(::MOI.Nonnegatives, T) = typemax(T) + +end # module diff --git a/ext/ExaModelsMOI.jl b/ext/ExaModelsMOI.jl index 6fc1d5a9..091220ba 100644 --- a/ext/ExaModelsMOI.jl +++ b/ext/ExaModelsMOI.jl @@ -43,7 +43,10 @@ function update_bin!(bin, e, p) if _update_bin!(bin, e, p) # if update succeeded, return the original bin return bin else # if update has failed, return a new bin - return Bin(e, [p], bin) + if p isa Tuple + p = [p] + end + return Bin(e, p, bin) end end function _update_bin!(bin::Bin{E,P,I}, e, p) where {E,P,I} @@ -61,6 +64,9 @@ end function check_supported(T, moim) con_types = MOI.get(moim, MOI.ListOfConstraintTypesPresent()) for (F, S) in con_types + if ExaModels.is_extension_type(F) + continue + end !(F <: SUPPORTED_FUNC_TYPE_WITH_VAR) && error("Unsupported function type $F.") if F <: MOI.VariableIndex !(S <: SUPPORTED_VAR_SET_TYPE) && @@ -71,7 +77,7 @@ function check_supported(T, moim) end obj_type = MOI.get(moim, MOI.ObjectiveFunctionType()) - !(obj_type <: SUPPORTED_FUNC_TYPE_WITH_VAR) && + !(obj_type <: SUPPORTED_FUNC_TYPE_WITH_VAR || ExaModels.is_extension_type(obj_type)) && error("Unsupported objective function type $obj_type.") obj_sense = MOI.get(moim, MOI.ObjectiveSense()) @@ -204,23 +210,19 @@ function copy_constraints!(c, moim, var_to_idx, T) con_types = MOI.get(moim, MOI.ListOfConstraintTypesPresent()) for (F, S) in con_types + F <: MOI.VariableIndex && continue + ExaModels.is_extension_type(F) && continue cis = MOI.get(moim, MOI.ListOfConstraintIndices{F,S}()) - if F <: MOI.VariableIndex - for ci in cis - vi = MOI.get(moim, MOI.ConstraintFunction(), ci) - vartype, var_idx = var_to_idx[vi] - if vartype === :variable - con_to_idx[ci] = var_idx - end - end - continue - end - bin, offset = - exafy_con(moim, cis, bin, offset, lcon, ucon, y0, var_to_idx, con_to_idx) + bin, offset = exafy_con(moim, cis, bin, offset, lcon, ucon, y0, var_to_idx, con_to_idx) end c, cons = ExaModels.add_con(c, offset; start = y0, lcon = lcon, ucon = ucon) c = build_constraint!(c, cons, bin) + # Hook for extensions (e.g. GenOpt) to add their constraint types + if applicable(ExaModels.copy_extra_constraints!, c, moim, var_to_idx, con_to_idx, T) + c = ExaModels.copy_extra_constraints!(c, moim, var_to_idx, con_to_idx, T) + end + return c, con_to_idx end @@ -323,7 +325,9 @@ function exafy_con( var_to_idx, con_to_idx, ) where {V<:Vector{<:MOI.ConstraintIndex}} - l = length(cons) + l = sum(cons) do ci + MOI.dimension(MOI.get(moim, MOI.ConstraintSet(), ci)) + end resize!(lcon, offset + l) resize!(ucon, offset + l) @@ -331,7 +335,7 @@ function exafy_con( for (i, ci) in enumerate(cons) func = MOI.get(moim, MOI.ConstraintFunction(), ci) set = MOI.get(moim, MOI.ConstraintSet(), ci) - con_to_idx[ci] = offset + i + con_to_idx[ci] = offset + 1 start = if MOI.supports( moim, MOI.ConstraintPrimalStart(), typeof(ci) ) @@ -342,8 +346,9 @@ function exafy_con( _exafy_con_update_start(ci, start, y0, con_to_idx) _exafy_con_update_vector(ci, set, lcon, ucon, con_to_idx) bin = _exafy_con(ci, func, bin, var_to_idx, con_to_idx) + offset += MOI.dimension(set) end - return bin, (offset += l) + return bin, offset end function _exafy_con_update_start(i, start, y0, con_to_idx) @@ -451,9 +456,12 @@ function exafy_obj(o::MOI.ScalarNonlinearFunction, bin, var_to_idx) bin = update_bin!(bin, e, p) end constant += m.constant - else + elseif m isa MOI.ScalarNonlinearFunction e, p = _exafy(m, var_to_idx) bin = update_bin!(bin, e, p) + else + e, p = ExaModels.exafy_extension_obj_arg(m, var_to_idx) + bin = update_bin!(bin, e, p) end end else @@ -464,6 +472,15 @@ function exafy_obj(o::MOI.ScalarNonlinearFunction, bin, var_to_idx) return update_bin!(bin, ExaModels.Null(constant), (1,)) # TODO see if this can be empty tuple end +# Fallback for extension objective types (e.g. SumGenerator as top-level objective) +function exafy_obj(o, bin, var_to_idx) + if !ExaModels.is_extension_type(typeof(o)) + throw(MOI.UnsupportedAttribute(MOI.ObjectiveFunction{typeof(o)}())) + end + e, p = ExaModels.exafy_extension_obj_arg(o, var_to_idx) + return update_bin!(bin, e, p) +end + function _exafy(v::MOI.VariableIndex, var_to_idx, p = ()) i = ExaModels.DataIndexed(ExaModels.DataSource(), length(p) + 1) vartype, idx = var_to_idx[v] @@ -481,7 +498,7 @@ function _exafy(i::R, var_to_idx, p) where {R<:Real} end function _exafy(e::MOI.ScalarNonlinearFunction, var_to_idx, p = ()) - return op(e.head)((begin + return ExaModels.op(e.head)((begin c, p = _exafy(e, var_to_idx, p) c end for e in e.args)...), p @@ -542,8 +559,7 @@ function _exafy(e::MOI.ScalarQuadraticTerm{T}, var_to_idx, p = ()) where {T} end end -# eval can be a performance killer -- we want to explicitly include symbols for frequently used operations. -function op(s::Symbol) +function ExaModels.op(s::Symbol) # uni/multi if s === :+ return + @@ -710,6 +726,14 @@ end MOI.is_empty(model::Optimizer) = isnothing(model.model) +function MOI.supports_constraint( + ::Optimizer, + ::Type{F}, + ::Type{S}, +) where {F<:MOI.AbstractFunction, S<:MOI.AbstractSet} + return ExaModels.is_extension_type(F) +end + function MOI.supports_constraint( ::Optimizer, ::Type{<:SUPPORTED_FUNC_TYPE}, @@ -730,6 +754,9 @@ end function MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{<:SUPPORTED_FUNC_TYPE_WITH_VAR}) return true end +function MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{F}) where {F} + return ExaModels.is_extension_type(F) +end function MOI.supports(::Optimizer, ::MOI.VariablePrimalStart, ::Type{MOI.VariableIndex}) return true end @@ -863,7 +890,7 @@ function _make_index_map(model::MOI.ModelLike, var_to_idx, con_to_idx) end end for (F, S) in MOI.get(model, MOI.ListOfConstraintTypesPresent()) - _make_constraints_map(model, map.con_map[F, S], con_to_idx) + _make_constraints_map(model, map.con_map[F, S], con_to_idx, var_to_idx) end return map end @@ -871,9 +898,27 @@ function _make_constraints_map( model, map::MOI.Utilities.DoubleDicts.IndexDoubleDictInner{F,S}, con_to_idx, + var_to_idx, ) where {F,S} for c in MOI.get(model, MOI.ListOfConstraintIndices{F,S}()) - map[c] = typeof(c)(con_to_idx[c]) + if haskey(con_to_idx, c) + map[c] = typeof(c)(con_to_idx[c]) + end + end + return +end +function _make_constraints_map( + model, + map::MOI.Utilities.DoubleDicts.IndexDoubleDictInner{MOI.VariableIndex,S}, + con_to_idx, + var_to_idx, +) where {S} + for c in MOI.get(model, MOI.ListOfConstraintIndices{MOI.VariableIndex,S}()) + vi = MOI.get(model, MOI.ConstraintFunction(), c) + entry = var_to_idx[vi] + if entry.type === :variable + map[c] = typeof(c)(entry.idx) + end end return end diff --git a/src/ExaModels.jl b/src/ExaModels.jl index 61ac4c85..587a7b63 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -60,6 +60,7 @@ include("deprecated.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..9ba1b846 --- /dev/null +++ b/src/wrapper.jl @@ -0,0 +1,35 @@ +# 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 + +""" + 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..f067cfac --- /dev/null +++ b/test/GenOptTest/GenOptTest.jl @@ -0,0 +1,156 @@ +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_exa_model(model) + be = backend(model) + return be.optimizer.model.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. + m = _get_exa_model(model) + @test m.meta.ncon == 36 + cons = m.cons + # Filter out the empty standard-constraint block + nonempty = filter(c -> length(c.itr) > 0, cons) + @test length(nonempty) == 10 + itr_lengths = sort([length(c.itr) for c in nonempty]) + @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 c4761037..11bf49b7 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -2,6 +2,7 @@ Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2" 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 f115393a..87df8d31 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -25,6 +25,7 @@ include("TwoStageTest/TwoStageTest.jl") include("GetterSetterTest/GetterSetterTest.jl") include("PrettyPrintTest.jl") # include("OptimalControlTest/OptimalControlTest.jl") +include("GenOptTest/GenOptTest.jl") @testset verbose = true "ExaModels test" begin @info "Running Deprecated API Test" @@ -57,6 +58,8 @@ include("PrettyPrintTest.jl") # @info "Running OptimalControl Test" # OptimalControlTest.runtests() + @info "Running GenOpt Test" + GenOptTest.runtests() end # Force full GC before Julia exits so that OpenCL/PoCL objects are finalized