Track variable bounds for ExaModels.Optimizer

ext/ExaModelsMOI.jl

@@ -195,7 +195,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
+        if F <: MOI.VariableIndex
+            cis = MOI.get(moim, MOI.ListOfConstraintIndices{F, S}())
+            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
+                else
+                    # error("Bound constraints on parameters are not supported")
+                end
+            end
+            continue
+        end
         cis = MOI.get(moim, MOI.ListOfConstraintIndices{F,S}())
         bin, offset =
             exafy_con(moim, cis, bin, offset, lcon, ucon, y0, var_to_idx, con_to_idx)
@@ -682,6 +694,8 @@ mutable struct Optimizer{B,S} <: MOI.AbstractOptimizer
     result::Any
     solve_time::Float64
     options::Dict{Symbol,Any}
+    var_to_idx::Dict{MOI.VariableIndex, NamedTuple{(:type, :idx), Tuple{Symbol, Int}}}
+    con_to_idx::Dict{MOI.ConstraintIndex, Int}
 end
 
 MOI.is_empty(model::Optimizer) = isnothing(model.model)
@@ -711,16 +725,30 @@ function MOI.supports(::Optimizer, ::MOI.VariablePrimalStart, ::Type{MOI.Variabl
 end
 
 function ExaModels.Optimizer(solver, backend = nothing; kwargs...)
-    return Optimizer(solver, backend, nothing, nothing, 0.0, Dict{Symbol,Any}(kwargs...))
+    return Optimizer(
+        solver,
+        backend,
+        nothing,
+        nothing,
+        0.0,
+        Dict{Symbol, Any}(kwargs...),
+        Dict{MOI.VariableIndex, NamedTuple{(:type, :idx), Tuple{Symbol, Int}}}(),
+        Dict{MOI.ConstraintIndex, Int}(),
+    )
 end
 
 function MOI.empty!(model::ExaModelsMOI.Optimizer)
     model.model = nothing
+    empty!(model.var_to_idx)
+    empty!(model.con_to_idx)
+    return
 end
 
 function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike)
     core, maps = to_exacore(src; backend = dest.backend)
     dest.model = ExaModels.ExaModel(core; prod = true)
+    dest.var_to_idx = maps[1]
+    dest.con_to_idx = maps[2]
     return _make_index_map(src, maps)
 end
 

test/JuMPTest/JuMPTest.jl

@@ -292,12 +292,21 @@ function runtests()
                     set_optimizer_attribute(jm, "print_level", 0)
                     optimize!(jm)
                     sol = value.(all_variables(jm))
+                    dsol = dual.(all_constraints(jm, include_variable_in_set_constraints = true))
+
+                    set_optimizer(jm, () -> ExaModels.Optimizer(ipopt))
+                    optimize!(jm)
+                    sol2 = value.(all_variables(jm))
+                    dsol2 = dual.(all_constraints(jm, include_variable_in_set_constraints = true))
+                    @test sol ≈ sol2 atol = 1.0e-6
+                    @test dsol ≈ dsol2 atol = 1.0e-6
 
                     for backend in BACKENDS
                         @testset "$backend" begin
                             m = WrapperNLPModel(ExaModel(jm; backend = backend))
                             result = ipopt(m; print_level = 0)
 
+                            # NOTE: assumes IndexMap is identity!
                             @test sol ≈ result.solution atol = 1e-6
                         end
                     end