Skip to content

Commit faf3389

Browse files
authored
Fix ordering of variables in QCQP.Optimizer (#113)
1 parent 0869d29 commit faf3389

File tree

7 files changed

+203
-47
lines changed

7 files changed

+203
-47
lines changed

src/QCQP/MOI_wrapper.jl

Lines changed: 26 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,24 @@ function _subs!(
228228
)
229229
end
230230

231+
"""
232+
_subs_ensure_moi_order(p::PolyJuMP.ScalarPolynomialFunction, old, new)
233+
234+
Substitutes old `MP.variables(p.polynomial)` with new vars, while re-sorting the
235+
MOI `p.variables` to get them in the correct order after substitution.
236+
"""
237+
function _subs_ensure_moi_order(p::PolyJuMP.ScalarPolynomialFunction, old, new)
238+
if isempty(old)
239+
return p
240+
end
241+
poly = MP.subs(p.polynomial, old => new)
242+
all_new_vars = MP.variables(poly)
243+
to_old_map = Dict(zip(new, old))
244+
to_moi_map = Dict(zip(MP.variables(p.polynomial), p.variables))
245+
moi_vars = [to_moi_map[get(to_old_map, v, v)] for v in all_new_vars]
246+
return PolyJuMP.ScalarPolynomialFunction(poly, moi_vars)
247+
end
248+
231249
function _subs!(
232250
p::PolyJuMP.ScalarPolynomialFunction,
233251
index_to_var::Dict{K,V},
@@ -244,10 +262,7 @@ function _subs!(
244262
index_to_var[vi] = var
245263
end
246264
end
247-
if !isempty(old_var)
248-
poly = MP.subs(p.polynomial, old_var => new_var)
249-
p = PolyJuMP.ScalarPolynomialFunction(poly, p.variables)
250-
end
265+
p = _subs_ensure_moi_order(p, old_var, new_var)
251266
return p, index_to_var
252267
end
253268

@@ -353,23 +368,13 @@ function MOI.Utilities.final_touch(model::Optimizer{T}, _) where {T}
353368
vars = _add_variables!(func, vars)
354369
monos = _add_monomials!(func, monos)
355370
end
356-
if !isempty(model.constraints)
357-
for S in keys(model.constraints)
358-
for ci in MOI.get(
359-
model,
360-
MOI.ListOfConstraintIndices{
361-
PolyJuMP.ScalarPolynomialFunction{
362-
T,
363-
model.constraints[S][1],
364-
},
365-
S,
366-
}(),
367-
)
368-
func = MOI.get(model, MOI.ConstraintFunction(), ci)
369-
func, index_to_var = _subs!(func, index_to_var)
370-
vars = _add_variables!(func, vars)
371-
monos = _add_monomials!(func, monos)
372-
end
371+
for (S, constraints) in model.constraints
372+
F = PolyJuMP.ScalarPolynomialFunction{T,constraints[1]}
373+
for ci in MOI.get(model, MOI.ListOfConstraintIndices{F,S}())
374+
func = MOI.get(model, MOI.ConstraintFunction(), ci)
375+
func, index_to_var = _subs!(func, index_to_var)
376+
vars = _add_variables!(func, vars)
377+
monos = _add_monomials!(func, monos)
373378
end
374379
end
375380
if !isnothing(monos)

src/functions.jl

Lines changed: 0 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -64,29 +64,6 @@ function Base.convert(
6464
return ScalarPolynomialFunction{T,P}(poly, variables)
6565
end
6666

67-
function _to_polynomial!(
68-
d::Dict{K,V},
69-
::Type{T},
70-
func::MOI.ScalarQuadraticFunction{T},
71-
) where {K,V,T}
72-
terms = MP.term_type(V, T)[MOI.constant(func)]
73-
for t in func.affine_terms
74-
push!(terms, MP.term(t.coefficient, _to_polynomial!(d, T, t.variable)))
75-
end
76-
for t in func.quadratic_terms
77-
coef = t.variable_1 == t.variable_2 ? t.coefficient / 2 : t.coefficient
78-
push!(
79-
terms,
80-
MP.term(
81-
coef,
82-
_to_polynomial!(d, T, t.variable_1) *
83-
_to_polynomial!(d, T, t.variable_2),
84-
),
85-
)
86-
end
87-
return MP.polynomial(terms)
88-
end
89-
9067
function Base.convert(
9168
::Type{ScalarPolynomialFunction{T,P}},
9269
func::MOI.ScalarQuadraticFunction{T},

src/nl_to_polynomial.jl

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,13 @@ function _to_polynomial!(
9595
elseif _is_operator(expr, :/) && length(operands) == 2
9696
a, b = _to_polynomial!.(Ref(d), T, operands)
9797
return a / b
98+
# TODO(odow): see PR#113
99+
# divisor, remainder = Base.divrem(a, b)
100+
# if iszero(remainder)
101+
# return divisor
102+
# else
103+
# return a / b
104+
# end
98105
elseif _is_variable(expr)
99106
return _to_polynomial!(d, T, operands[1])
100107
else
@@ -109,7 +116,8 @@ function _to_polynomial(expr, ::Type{T}) where {T}
109116
end
110117

111118
function _scalar_polynomial(d::Dict{K,V}, ::Type{T}, poly) where {T,K,V}
112-
variable_map = collect(d)
119+
var_set = Set(MP.variables(poly))
120+
variable_map = Tuple{K,V}[(k, v) for (k, v) in d if v in var_set]
113121
sort!(variable_map, by = x -> x[2], rev = true)
114122
variables = [x[1] for x in variable_map]
115123
P = MP.polynomial_type(V, T)

test/Project.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ MultivariateMoments = "f4abf1af-0426-5881-a0da-e2f168889b5e"
1010
MultivariatePolynomials = "102ac46a-7ee4-5c85-9060-abc95bfdeaa3"
1111
MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
1212
PolyJuMP = "ddf597a6-d67e-5340-b84c-e37d84115374"
13+
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
1314
SemialgebraicSets = "8e049039-38e8-557d-ae3a-bc521ccf6204"
1415
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
1516
TypedPolynomials = "afbbf031-7a57-5f58-a1b9-b774a0fad08d"

test/qcqp.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -149,8 +149,8 @@ function test_no_monomials(x, y, T)
149149
model = PolyJuMP.JuMP.GenericModel{T}() do
150150
return PolyJuMP.QCQP.Optimizer{T}(MOI.Utilities.MockOptimizer(inner))
151151
end
152-
PolyJuMP.@variable(model, 0 <= x[1:2] <= 1)
153-
PolyJuMP.@constraint(model, x[1] * x[2] == T(1))
152+
PolyJuMP.@variable(model, 0 <= x[1:2] <= 2)
153+
PolyJuMP.@constraint(model, x[1] * x[2] == 1)
154154
PolyJuMP.@objective(model, Min, sum(x))
155155
PolyJuMP.optimize!(model)
156156
@test MOI.get(inner, MOI.NumberOfVariables()) == 2

test/qcqp_extra.jl

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
module TestQCQPExtra
2+
3+
using Test
4+
5+
import MathOptInterface as MOI
6+
import MultivariatePolynomials as MP
7+
import PolyJuMP
8+
import Random
9+
10+
MOI.Utilities.@model(
11+
Model,
12+
(),
13+
(MOI.LessThan, MOI.GreaterThan, MOI.EqualTo, MOI.Interval),
14+
(),
15+
(),
16+
(),
17+
(MOI.ScalarAffineFunction, MOI.ScalarQuadraticFunction),
18+
(),
19+
(),
20+
)
21+
22+
function MOI.supports(
23+
::Model,
24+
::MOI.ObjectiveFunction{MOI.ScalarNonlinearFunction},
25+
)
26+
return false
27+
end
28+
29+
function run_tests_e2e()
30+
for name in names(@__MODULE__; all = true)
31+
if startswith("$name", "test_e2e")
32+
@testset "$(name) $T" for T in [Int, Float64]
33+
getfield(@__MODULE__, name)(T)
34+
end
35+
end
36+
end
37+
return
38+
end
39+
40+
function run_test_scalar_polynomial_function(xs, samples)
41+
@testset "qcqp extra $T" for T in [Float64, BigFloat]
42+
for i in 1:samples
43+
test_scalar_polynomial_function(xs, T)
44+
end
45+
end
46+
return
47+
end
48+
49+
function run_tests_subs(xs, ys, samples, TS)
50+
Random.seed!(2024)
51+
for name in names(@__MODULE__; all = true)
52+
if startswith("$name", "test_subs")
53+
@testset "$(name) $T" for T in TS
54+
for i in 1:samples
55+
getfield(@__MODULE__, name)(xs, ys, T)
56+
end
57+
end
58+
end
59+
end
60+
return
61+
end
62+
63+
function test_unconstrained_before_projection(T)
64+
inner = Model{T}()
65+
optimizer = MOI.Utilities.MockOptimizer(inner)
66+
model = PolyJuMP.JuMP.GenericModel{T}(
67+
() -> PolyJuMP.QCQP.Optimizer{T}(optimizer),
68+
)
69+
PolyJuMP.@variable(model, -1 <= a[1:2] <= 1)
70+
PolyJuMP.@objective(model, Min, a[1]^2 * a[2]^2)
71+
PolyJuMP.optimize!(model)
72+
vis = MOI.get(inner, MOI.ListOfVariableIndices())
73+
@test length(vis) == 4
74+
F = MOI.ScalarQuadraticFunction{T}
75+
S = MOI.EqualTo{T}
76+
cis = MOI.get(inner, MOI.ListOfConstraintIndices{F,S}())
77+
@test length(cis) == 2
78+
return
79+
end
80+
81+
function test_unconstrained_after_projection(T)
82+
inner = Model{T}()
83+
optimizer = MOI.Utilities.MockOptimizer(inner)
84+
model = PolyJuMP.JuMP.GenericModel{T}(
85+
() -> PolyJuMP.QCQP.Optimizer{T}(optimizer),
86+
)
87+
PolyJuMP.@variable(model, -1 <= a <= 1)
88+
PolyJuMP.@objective(model, Min, a^2)
89+
PolyJuMP.optimize!(model)
90+
vis = MOI.get(inner, MOI.ListOfVariableIndices())
91+
@test length(vis) == 1
92+
F = MOI.ScalarQuadraticFunction{T}
93+
S = MOI.EqualTo{T}
94+
cis = MOI.get(inner, MOI.ListOfConstraintIndices{F,S}())
95+
@test length(cis) == 0
96+
return
97+
end
98+
99+
function _random_polynomial(vars, T)
100+
ms = Random.shuffle(MP.monomials(vars, 1:length(vars)))
101+
return sum(ms[i] * T(randn()) for i in eachindex(ms) if rand(Bool))
102+
end
103+
104+
function test_subs!_preserves_moi_sync(xs, ys, T)
105+
p = _random_polynomial(xs, T)
106+
mois = MOI.VariableIndex.(eachindex(xs))
107+
vals = T.(randn(length(mois)))
108+
mask = rand(Bool, length(xs))
109+
is = Random.shuffle(eachindex(xs)[mask])
110+
index = Dict{eltype(mois),eltype(xs)}(zip(mois[is], ys[is]))
111+
moi_map = Dict(zip(xs, mois))
112+
moivars = [moi_map[v] for v in MP.variables(p)]
113+
before = PolyJuMP.ScalarPolynomialFunction(p, moivars)
114+
after, _ = PolyJuMP.QCQP._subs!(before, index)
115+
bmap = [vals[v.value] for v in before.variables]
116+
amap = [vals[v.value] for v in after.variables]
117+
bvalue = before.polynomial(MP.variables(before.polynomial) => bmap)
118+
avalue = after.polynomial(MP.variables(after.polynomial) => amap)
119+
# avoid verbose fails
120+
@test isapprox(Float64(bvalue), Float64(avalue))
121+
return
122+
end
123+
124+
function test_scalar_polynomial_function(xs, T)
125+
pick = rand(eachindex(xs))
126+
ids = Random.shuffle(eachindex(xs))
127+
poly = sum(T(randn()) * xs[i] for i in ids if i != pick)
128+
mois = MOI.VariableIndex.(eachindex(xs))
129+
moi_to_vars = Dict(zip(mois, xs))
130+
spf = PolyJuMP._scalar_polynomial(moi_to_vars, Any, poly)
131+
expected = MP.variables(poly)
132+
actual = [moi_to_vars[m] for m in spf.variables]
133+
@test length(MP.variables(spf.polynomial)) == length(spf.variables)
134+
@test expected == actual
135+
return
136+
end
137+
138+
end # module
139+
140+
using Test
141+
142+
@testset "TestQCQPFinalTouch" begin
143+
TestQCQPExtra.run_tests_e2e()
144+
end
145+
146+
import DynamicPolynomials
147+
@testset "DynamicPolynomials" begin
148+
ids = 1:4
149+
DynamicPolynomials.@polyvar(x[ids])
150+
DynamicPolynomials.@polyvar(y[ids])
151+
samples = 10
152+
types = [Float64, BigFloat] # Rational fails with DynamicPolynomials
153+
TestQCQPExtra.run_tests_subs(x, y, samples, types)
154+
TestQCQPExtra.run_test_scalar_polynomial_function(x, samples)
155+
end
156+
157+
import TypedPolynomials
158+
@testset "TypedPolynomials" begin
159+
TypedPolynomials.@polyvar(z[1:4])
160+
TypedPolynomials.@polyvar(w[1:4])
161+
types = [Float64, BigFloat, Rational{BigInt}]
162+
samples = 10
163+
TestQCQPExtra.run_tests_subs(z, w, samples, types)
164+
end

test/runtests.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,3 +8,4 @@ include("Mock/mock_tests.jl")
88
include("kkt.jl")
99
include("sage.jl")
1010
include("qcqp.jl")
11+
include("qcqp_extra.jl")

0 commit comments

Comments
 (0)