Add extension to GenOpt

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"

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

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;

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,

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