Add extension to GenOpt

Author: blegat

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"
 ExaModelsIpopt = ["MathOptInterface", "NLPModelsIpopt"]
 ExaModelsJuMP = "JuMP"
 ExaModelsKernelAbstractions = "KernelAbstractions"
@@ -36,6 +38,7 @@ ExaModelsSpecialFunctions = "SpecialFunctions"
 
 [compat]
 Adapt = "4"
+GenOpt = "0.1"
 Ipopt = "1.11"
 JuMP = "1"
 KernelAbstractions = "0.9"

ext/ExaModelsGenOpt.jl

@@ -0,0 +1,93 @@
+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)
+    return _exagen(m.func, m.iterators)
+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}())
+        _copy_generator_constraints!(c, moim, cis, var_to_idx, con_to_idx, T, S)
+    end
+end
+
+function _copy_generator_constraints!(c, moim, cis, var_to_idx, con_to_idx, T, ::Type{S}) where {S}
+    # FIXME we assume that `var_to_idx` is the identity
+    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)
+        ExaModels.constraint(c, expr, pars; lcon = _lower_bounds(set, T), ucon = _upper_bounds(set, T))
+    end
+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.ParSource()
+            else
+                return ExaModels.ParIndexed(ExaModels.ParSource(), 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) for e in f.args)...)
+    end
+end
+
+function _exagen(func::MOI.ScalarNonlinearFunction, iterators)
+    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)
+    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