Skip to content

Commit

Permalink
Fix ordering of variables in QCQP.Optimizer (#113)
Browse files Browse the repository at this point in the history
  • Loading branch information
votroto authored Mar 17, 2024
1 parent 0869d29 commit faf3389
Show file tree
Hide file tree
Showing 7 changed files with 203 additions and 47 deletions.
47 changes: 26 additions & 21 deletions src/QCQP/MOI_wrapper.jl
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,24 @@ 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},
Expand All @@ -244,10 +262,7 @@ function _subs!(
index_to_var[vi] = var
end
end
if !isempty(old_var)
poly = MP.subs(p.polynomial, old_var => new_var)
p = PolyJuMP.ScalarPolynomialFunction(poly, p.variables)
end
p = _subs_ensure_moi_order(p, old_var, new_var)
return p, index_to_var
end

Expand Down Expand Up @@ -353,23 +368,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, constraints) in model.constraints
F = PolyJuMP.ScalarPolynomialFunction{T,constraints[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)
Expand Down
23 changes: 0 additions & 23 deletions src/functions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -64,29 +64,6 @@ function Base.convert(
return ScalarPolynomialFunction{T,P}(poly, variables)
end

function _to_polynomial!(
d::Dict{K,V},
::Type{T},
func::MOI.ScalarQuadraticFunction{T},
) where {K,V,T}
terms = MP.term_type(V, T)[MOI.constant(func)]
for t in func.affine_terms
push!(terms, MP.term(t.coefficient, _to_polynomial!(d, T, t.variable)))
end
for t in func.quadratic_terms
coef = t.variable_1 == t.variable_2 ? t.coefficient / 2 : t.coefficient
push!(
terms,
MP.term(
coef,
_to_polynomial!(d, T, t.variable_1) *
_to_polynomial!(d, T, t.variable_2),
),
)
end
return MP.polynomial(terms)
end

function Base.convert(
::Type{ScalarPolynomialFunction{T,P}},
func::MOI.ScalarQuadraticFunction{T},
Expand Down
10 changes: 9 additions & 1 deletion src/nl_to_polynomial.jl
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,13 @@ function _to_polynomial!(
elseif _is_operator(expr, :/) && length(operands) == 2
a, b = _to_polynomial!.(Ref(d), T, operands)
return a / b
# TODO(odow): see PR#113
# divisor, remainder = Base.divrem(a, b)
# if iszero(remainder)
# return divisor
# else
# return a / b
# end
elseif _is_variable(expr)
return _to_polynomial!(d, T, operands[1])
else
Expand All @@ -109,7 +116,8 @@ function _to_polynomial(expr, ::Type{T}) where {T}
end

function _scalar_polynomial(d::Dict{K,V}, ::Type{T}, poly) where {T,K,V}
variable_map = collect(d)
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)
Expand Down
1 change: 1 addition & 0 deletions test/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ MultivariateMoments = "f4abf1af-0426-5881-a0da-e2f168889b5e"
MultivariatePolynomials = "102ac46a-7ee4-5c85-9060-abc95bfdeaa3"
MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
PolyJuMP = "ddf597a6-d67e-5340-b84c-e37d84115374"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
SemialgebraicSets = "8e049039-38e8-557d-ae3a-bc521ccf6204"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
TypedPolynomials = "afbbf031-7a57-5f58-a1b9-b774a0fad08d"
4 changes: 2 additions & 2 deletions test/qcqp.jl
Original file line number Diff line number Diff line change
Expand Up @@ -149,8 +149,8 @@ function test_no_monomials(x, y, T)
model = PolyJuMP.JuMP.GenericModel{T}() do
return PolyJuMP.QCQP.Optimizer{T}(MOI.Utilities.MockOptimizer(inner))
end
PolyJuMP.@variable(model, 0 <= x[1:2] <= 1)
PolyJuMP.@constraint(model, x[1] * x[2] == T(1))
PolyJuMP.@variable(model, 0 <= x[1:2] <= 2)
PolyJuMP.@constraint(model, x[1] * x[2] == 1)
PolyJuMP.@objective(model, Min, sum(x))
PolyJuMP.optimize!(model)
@test MOI.get(inner, MOI.NumberOfVariables()) == 2
Expand Down
164 changes: 164 additions & 0 deletions test/qcqp_extra.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,164 @@
module TestQCQPExtra

using Test

import MathOptInterface as MOI
import MultivariatePolynomials as MP
import PolyJuMP
import Random

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
return
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
return
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
return
end

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
return
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
return
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))
return
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
return
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
1 change: 1 addition & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -8,3 +8,4 @@ include("Mock/mock_tests.jl")
include("kkt.jl")
include("sage.jl")
include("qcqp.jl")
include("qcqp_extra.jl")

0 comments on commit faf3389

Please sign in to comment.