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"
 JuMP = "1"
 KernelAbstractions = "0.9"
 MathOptInterface = "1.19"

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

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,15 +325,17 @@ 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)
     resize!(y0, offset + l)
     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
@@ -884,17 +911,35 @@ 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
 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

src/ExaModels.jl

@@ -61,6 +61,7 @@ include("deprecated.jl")
 include("utils.jl")
 include("tags.jl")
 include("two_stage.jl")
+include("wrapper.jl")
 
 export ExaModel,
     ExaCore,

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