Add tests for non-continguous variable indices

Author: blegat

ext/ExaModelsGenOpt.jl

@@ -10,8 +10,8 @@ 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)
+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
@@ -26,32 +26,34 @@ function ExaModels.copy_extra_constraints!(c, moim, var_to_idx, con_to_idx, T)
 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)
+        expr, pars = _exagen(func.func, func.iterators, var_to_idx)
         c, _ = ExaModels.add_con(c, expr, pars; lcon = _lower_bounds(set, T), ucon = _upper_bounds(set, T))
     end
     return c
 end
 
 # Convert GenOpt expression trees to ExaModels format
 
-exagen(α::Number, _) = α
+exagen(α::Number, _, _) = α
 
-function exagen(f::MOI.ScalarNonlinearFunction, offsets)
+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)
-            if !iszero(v.offset)
-                idx = v.offset + idx
+            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) - 1)
+                idx += cp[i-2] * (exagen(f.args[i], offsets, var_to_idx) - 1)
             end
             return ExaModels.Var(idx)
         elseif v isa IteratorIndex
@@ -67,11 +69,11 @@ function exagen(f::MOI.ScalarNonlinearFunction, offsets)
             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)...)
+        return ExaModels.op(f.head)((exagen(e, offsets, var_to_idx) for e in f.args)...)
     end
 end
 
-function _exagen(func::MOI.ScalarNonlinearFunction, iterators)
+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
@@ -82,7 +84,7 @@ function _exagen(func::MOI.ScalarNonlinearFunction, iterators)
             reduce((i, j) -> tuple(i..., j...), I)
         end)
     end
-    expr = exagen(func, cs)
+    expr = exagen(func, cs, var_to_idx)
     return expr, pars
 end
 

ext/ExaModelsMOI.jl

@@ -458,7 +458,7 @@ function exafy_obj(o::MOI.ScalarNonlinearFunction, bin, var_to_idx)
                 constant += m.constant
             else
                 # Try extension hook first (e.g. for SumGenerator)
-                result = ExaModels.exafy_extension_obj_arg(m)
+                result = ExaModels.exafy_extension_obj_arg(m, var_to_idx)
                 if !isnothing(result)
                     e, p = result
                     bin = update_bin!(bin, e, p)

src/wrapper.jl

@@ -18,13 +18,14 @@ function is_extension_type end
 is_extension_type(::Type) = false
 
 """
-    exafy_extension_obj_arg(m) -> Union{Nothing, Tuple}
+    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. Returns `nothing` if the type is not handled by any extension.
+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
-exafy_extension_obj_arg(m) = nothing
+exafy_extension_obj_arg(m, var_to_idx) = nothing
 
 """
     op(s::Symbol)

test/GenOptTest/GenOptTest.jl

@@ -53,6 +53,60 @@ function _get_exa_model(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()
     @testset "GenOpt extension test" begin
         @testset "Quadrotor with ExaModels.Optimizer" begin
@@ -81,6 +135,19 @@ function runtests()
             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 = 1e-3)
+        end
     end
 end