add tests + refactor

src/QCQP/MOI_wrapper.jl

@@ -211,6 +211,28 @@ function _subs!(
     )
 end
 
+"""
+    _subs_ensure_moi_order(p::PolyJuMP.ScalarPolynomialFunction, old, new)
+
+Substitutes old `MP.variables(p.polynomial)` with new vars, while re-sorting the
+MOI `p.variables` to get them in the correct order after substitution.
+"""
+function _subs_ensure_moi_order(p::PolyJuMP.ScalarPolynomialFunction, old, new)
+    if isempty(old)
+        return p
+    end
+
+    poly = MP.subs(p.polynomial, old => new)
+
+    all_new_vars = MP.variables(poly)
+    to_old_map = Dict(zip(new, old))
+    to_moi_map = Dict(zip(MP.variables(p.polynomial), p.variables))
+    moi_vars = [to_moi_map[get(to_old_map, v, v)] for v in all_new_vars]
+
+    return PolyJuMP.ScalarPolynomialFunction(poly, moi_vars)
+end
+
+
 function _subs!(
     p::PolyJuMP.ScalarPolynomialFunction,
     index_to_var::Dict{K,V},
@@ -227,16 +249,7 @@ function _subs!(
             index_to_var[vi] = var
         end
     end
-    if !isempty(old_var)
-        to_old_map = Dict(zip(new_var, old_var))
-        to_moi_map = Dict(zip(MP.variables(p.polynomial), p.variables))
-
-        poly = MP.subs(p.polynomial, old_var => new_var)
-        all_new_vars = MP.variables(poly)
-        moi_vars = [to_moi_map[get(to_old_map, v, v)] for v in all_new_vars]
-
-        p = PolyJuMP.ScalarPolynomialFunction(poly, moi_vars)
-    end
+    p = _subs_ensure_moi_order(p, old_var, new_var)
     return p, index_to_var
 end
 
@@ -342,23 +355,13 @@ function MOI.Utilities.final_touch(model::Optimizer{T}, _) where {T}
         vars = _add_variables!(func, vars)
         monos = _add_monomials!(func, monos)
     end
-    if !isempty(model.constraints)
-        for S in keys(model.constraints)
-            for ci in MOI.get(
-                model,
-                MOI.ListOfConstraintIndices{
-                    PolyJuMP.ScalarPolynomialFunction{
-                        T,
-                        model.constraints[S][1],
-                    },
-                    S,
-                }(),
-            )
-                func = MOI.get(model, MOI.ConstraintFunction(), ci)
-                func, index_to_var = _subs!(func, index_to_var)
-                vars = _add_variables!(func, vars)
-                monos = _add_monomials!(func, monos)
-            end
+    for S in keys(model.constraints)
+        F = PolyJuMP.ScalarPolynomialFunction{T,model.constraints[S][1]}
+        for ci in MOI.get(model, MOI.ListOfConstraintIndices{F,S}())
+            func = MOI.get(model, MOI.ConstraintFunction(), ci)
+            func, index_to_var = _subs!(func, index_to_var)
+            vars = _add_variables!(func, vars)
+            monos = _add_monomials!(func, monos)
         end
     end
     if !isnothing(monos)

src/nl_to_polynomial.jl

@@ -114,8 +114,10 @@ function _to_polynomial(expr, ::Type{T}) where {T}
 end
 
 function _scalar_polynomial(d::Dict{K,V}, ::Type{T}, poly) where {T,K,V}
-    inv = Dict(v => k for (k, v) in d)
-    variables = [inv[v] for v in MP.variables(poly)]
+    var_set = Set(MP.variables(poly))
+    variable_map = Tuple{K, V}[(k, v) for (k, v) in d if v in var_set]
+    sort!(variable_map, by = x -> x[2], rev = true)
+    variables = [x[1] for x in variable_map]
     P = MP.polynomial_type(V, T)
     return ScalarPolynomialFunction{T,P}(poly, variables)
 end

test/qcqp_extra.jl

@@ -0,0 +1,174 @@
+module TestQCQPExtra
+
+using Test
+
+import MathOptInterface as MOI
+import MultivariatePolynomials as MP
+import PolyJuMP
+import Random
+
+# Random.seed! set below!
+
+function test_unconstrained_before_projection(T)
+    inner = Model{T}()
+    optimizer = MOI.Utilities.MockOptimizer(inner)
+    model = PolyJuMP.JuMP.GenericModel{T}(
+        () -> PolyJuMP.QCQP.Optimizer{T}(optimizer),
+    )
+    PolyJuMP.@variable(model, -1 <= a[1:2] <= 1)
+    PolyJuMP.@objective(model, Min, a[1]^2*a[2]^2)
+    PolyJuMP.optimize!(model)
+
+    vis = MOI.get(inner, MOI.ListOfVariableIndices())
+    @test length(vis) == 4
+
+    F = MOI.ScalarQuadraticFunction{T}
+    S = MOI.EqualTo{T}
+    cis = MOI.get(inner, MOI.ListOfConstraintIndices{F,S}())
+    @test length(cis) == 2
+end
+
+function test_unconstrained_after_projection(T)
+    inner = Model{T}()
+    optimizer = MOI.Utilities.MockOptimizer(inner)
+    model = PolyJuMP.JuMP.GenericModel{T}(
+        () -> PolyJuMP.QCQP.Optimizer{T}(optimizer),
+    )
+    PolyJuMP.@variable(model, -1 <= a <= 1)
+    PolyJuMP.@objective(model, Min, a^2)    
+    PolyJuMP.optimize!(model)
+
+    vis = MOI.get(inner, MOI.ListOfVariableIndices())
+    @test length(vis) == 1
+
+    F = MOI.ScalarQuadraticFunction{T}
+    S = MOI.EqualTo{T}
+    cis = MOI.get(inner, MOI.ListOfConstraintIndices{F,S}())
+    @test length(cis) == 0
+end
+
+function _random_polynomial(vars, T)
+    ms = Random.shuffle(MP.monomials(vars, 1:length(vars)))
+    return sum(ms[i] * T(randn()) for i in eachindex(ms) if rand(Bool))
+end
+
+function test_subs!_preserves_moi_sync(xs, ys, T)
+    p = _random_polynomial(xs, T)
+
+    mois = MOI.VariableIndex.(eachindex(xs))
+    vals = T.(randn(length(mois)))
+
+    mask = rand(Bool, length(xs))
+    is = Random.shuffle(eachindex(xs)[mask])
+    index = Dict{eltype(mois), eltype(xs)}(zip(mois[is], ys[is]))
+
+    moi_map = Dict(zip(xs, mois))
+    moivars = [moi_map[v] for v in MP.variables(p)]
+
+    before = PolyJuMP.ScalarPolynomialFunction(p, moivars)
+    after, _ = PolyJuMP.QCQP._subs!(before, index)
+
+    bmap = [vals[v.value] for v in before.variables]
+    amap = [vals[v.value] for v in after.variables]
+
+    bvalue = before.polynomial(MP.variables(before.polynomial) => bmap)
+    avalue = after.polynomial(MP.variables(after.polynomial) => amap)
+
+    # avoid verbose fails
+    @test isapprox(Float64(bvalue), Float64(avalue))
+end
+
+function test_scalar_polynomial_function(xs, T)
+    pick = rand(eachindex(xs))
+    ids = Random.shuffle(eachindex(xs))
+    poly = sum(T(randn()) * xs[i] for i in ids if i != pick)
+    mois = MOI.VariableIndex.(eachindex(xs))
+    moi_to_vars = Dict(zip(mois, xs))
+
+    spf = PolyJuMP._scalar_polynomial(moi_to_vars, Any, poly)
+
+    expected = MP.variables(poly)
+    actual = [moi_to_vars[m] for m in spf.variables]
+
+    @test length(MP.variables(spf.polynomial)) == length(spf.variables)
+    @test expected == actual
+end
+
+MOI.Utilities.@model(
+    Model,
+    (),
+    (MOI.LessThan, MOI.GreaterThan, MOI.EqualTo, MOI.Interval),
+    (),
+    (),
+    (),
+    (MOI.ScalarAffineFunction, MOI.ScalarQuadraticFunction),
+    (),
+    (),
+)
+
+function MOI.supports(
+    ::Model,
+    ::MOI.ObjectiveFunction{MOI.ScalarNonlinearFunction},
+)
+    return false
+end
+
+function run_tests_e2e()
+    for name in names(@__MODULE__; all = true)
+        if startswith("$name", "test_e2e")
+            @testset "$(name) $T" for T in [Int, Float64]
+                getfield(@__MODULE__, name)(T)
+            end
+        end
+    end
+end
+
+function run_test_scalar_polynomial_function(xs, samples)
+    @testset "qcqp extra $T" for T in [Float64, BigFloat]
+        for i in 1:samples
+            test_scalar_polynomial_function(xs, T)
+        end
+    end
+end
+
+function run_tests_subs(xs, ys, samples, TS)
+    Random.seed!(2024)
+    for name in names(@__MODULE__; all = true)
+        if startswith("$name", "test_subs")
+            @testset "$(name) $T" for T in TS
+                for i in 1:samples
+                    getfield(@__MODULE__, name)(xs, ys, T)
+                end
+            end
+        end
+    end
+end
+
+end # module
+
+using Test
+
+@testset "TestQCQPFinalTouch" begin
+    TestQCQPExtra.run_tests_e2e()
+end
+
+import DynamicPolynomials
+@testset "DynamicPolynomials" begin
+    ids = 1:4
+    DynamicPolynomials.@polyvar(x[ids])
+    DynamicPolynomials.@polyvar(y[ids])
+    samples = 10
+    types = [Float64, BigFloat] # Rational fails with DynamicPolynomials
+    TestQCQPExtra.run_tests_subs(x, y, samples, types)
+
+    TestQCQPExtra.run_test_scalar_polynomial_function(x, samples)
+end
+
+import TypedPolynomials
+@testset "TypedPolynomials" begin
+    TypedPolynomials.@polyvar(z[1:4])
+    TypedPolynomials.@polyvar(w[1:4])
+    types = [Float64, BigFloat, Rational{BigInt}]
+    samples = 10
+    TestQCQPExtra.run_tests_subs(z, w, samples, types)
+end