From b6debdcb522abe114c092257b23531b765d2384d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Sun, 6 Mar 2022 11:52:26 -0500 Subject: [PATCH] Add SAGE extension in doc --- .../tutorials/Extension/relative_entropy.jl | 227 ++++++++++++++++++ 1 file changed, 227 insertions(+) create mode 100644 docs/src/tutorials/Extension/relative_entropy.jl diff --git a/docs/src/tutorials/Extension/relative_entropy.jl b/docs/src/tutorials/Extension/relative_entropy.jl new file mode 100644 index 000000000..24c6ed0da --- /dev/null +++ b/docs/src/tutorials/Extension/relative_entropy.jl @@ -0,0 +1,227 @@ +# # Relative Entropy + +#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Extension/relative_entropy.ipynb) +#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Extension/relative_entropy.ipynb) +# **Contributed by**: Benoît Legat + +using Test #src +using DynamicPolynomials +@polyvar x y +motzkin = x^4*y^2 + x^2*y^4 + 1 - 3x^2*y^2 + +module RelativeEntropy + +import MutableArithmetics +const MA = MutableArithmetics +import MultivariatePolynomials +const MP = MultivariatePolynomials +import MathOptInterface +const MOI = MathOptInterface +import JuMP +import PolyJuMP + +# **S**ums of **A**M/**G**M **E**xponential +struct SAGECone <: MOI.AbstractVectorSet + α::Matrix{Int} +end +MOI.dimension(set::SAGECone) = size(set.α, 1) +Base.copy(set::SAGECone) = set + +struct SAGESet <: PolyJuMP.PolynomialSet end +JuMP.reshape_set(::SAGECone, ::PolyJuMP.PolynomialShape) = SAGESet() +function _exponents_matrix(monos) + α = Matrix{Int}(undef, length(monos), MP.nvariables(monos)) + for (i, mono) in enumerate(monos) + α[i, :] = MP.exponents(mono) + end + return α +end +JuMP.moi_set(set::SAGESet, monos) = SAGECone(_exponents_matrix(monos)) + +struct AGECone <: MOI.AbstractVectorSet + α::Matrix{Int} + k::Int +end +MOI.dimension(set::AGECone) = size(set.α, 1) +Base.copy(set::AGECone) = set + +struct AGESet{MT <: MP.AbstractMonomial} <: PolyJuMP.PolynomialSet + monomial::MT +end +function JuMP.reshape_set(set::AGECone, shape::PolyJuMP.PolynomialShape) + return AGESet(shape.monomials[set.k]) +end +function JuMP.moi_set(set::AGESet, monos) + k = findfirst(isequal(set.monomial), monos) + return AGECone(_exponents_matrix(monos), k) +end + +function setdefaults!(data::PolyJuMP.Data) + PolyJuMP.setdefault!(data, PolyJuMP.NonNegPoly, SAGESet) + return +end + +function JuMP.build_constraint(_error::Function, p, set::Union{SAGESet,AGESet}; kws...) + coefs = PolyJuMP.non_constant_coefficients(p) + monos = MP.monomials(p) + cone = JuMP.moi_set(set, monos) + shape = PolyJuMP.PolynomialShape(monos) + return PolyJuMP.bridgeable( + JuMP.VectorConstraint(coefs, cone, shape), + JuMP.moi_function_type(typeof(coefs)), + typeof(cone), + ) +end + +struct SAGEBridge{T,F,G} <: MOI.Bridges.Constraint.AbstractBridge + ν::Matrix{MOI.VariableIndex} + age_constraints::Vector{MOI.ConstraintIndex{MOI.VectorOfVariables,AGECone}} + equality_constraints::Vector{MOI.ConstraintIndex{F,MOI.EqualTo{T}}} +end + +function MOI.Bridges.Constraint.bridge_constraint( + ::Type{SAGEBridge{T,F,G}}, + model, + func::G, + set::SAGECone, +) where {T,F,G} + m = size(set.α, 1) + ν = Vector{Vector{MOI.VariableIndex}}(undef, m) + age_constraints = Vector{MOI.ConstraintIndex{MOI.VectorOfVariables,AGECone}}(undef, m) + for k in 1:m + ν[i], age_constraints[i] = MOI.add_constrained_variables(model, AGECone(set.α, k)) + end + scalars = MOI.Utilities.eachscalar(func) + n = size(set.α, 2) + equality_constraints = map(1:n) do i + f = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(one(T), ν[:, i]), zero(T)) + MOI.add_constraint(model, MA.sub!(f, scalars[i]), MOI.EqualTo(zero(T))) + end + return SAGEBridge{T,F,G}(ν, age_constraints, equality_constraints) +end + +function MOI.supports_constraint( + ::Type{<:SAGEBridge{T}}, + ::Type{<:MOI.AbstractVectorFunction}, + ::Type{<:SAGECone}, +) where {T} + return true +end + +function MOI.Bridges.added_constrained_variable_types(::Type{<:SAGEBridge}) + return Tuple{Type}[AGECone] +end + +function MOI.Bridges.added_constraint_types(::Type{<:SAGEBridge{T,F}}) where {T,F} + return Tuple{Type,Type}[(F, MOI.EqualTo{T})] +end + +function MOI.Bridges.Constraint.concrete_bridge_type( + ::Type{<:SAGEBridge{T}}, + G::Type{<:MOI.AbstractVectorFunction}, + ::Type{<:SAGECone}, +) where {T} + S = MOI.Utilities.scalar_type(G) + F = MOI.Utilities.promote_operation(-, T, MOI.ScalarAffineFunction{Float64}, S) + return AGEBridge{T,F,G} +end + +function PolyJuMP.bridges( + ::Type{<:MOI.AbstractVectorFunction}, + ::Type{<:SAGECone}, +) + return [SAGEBridge] +end + +struct AGEBridge{T,F,G,H} <: MOI.Bridges.Constraint.AbstractBridge + k::Int + equality_constraints::Vector{MOI.ConstraintIndex{F,MOI.EqualTo{T}}} + relative_entropy_constraint::MOI.ConstraintIndex{G,MOI.RelativeEntropyCone} +end # See https://arxiv.org/pdf/1907.00814.pdf + +function MOI.Bridges.Constraint.bridge_constraint( + ::Type{AGEBridge{T,F,G,H}}, + model, + func::G, + set::AGECone, +) where {T,F,G,H} + m = size(set.α, 1) + ν = MOI.add_variables(model, m - 1) + sumν = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(one(T), ν), zero(T)) + ceq = map(1:size(set.α, 2)) do var + f = -sumν + j = 0 + for i in 1:m + if i != set.k + j += 1 + MA.add_mul!(f, convert(T, set.α[i, var]), MOI.SingleVariable(ν[j])) + end + end + MOI.Utilities.normalize_and_add_constraint(model, f, MOI.EqualTo(zero(T))) + end + scalars = MOI.Utilities.eachscalar(func) + f = MOI.Utilities.operate( + vcat, + T, + scalars[set.k] + sumν, + scalars[setdiff(1:m, set.k)], + MOI.VectorOfVariables(ν) + ) + relative_entropy_constraint = MOI.add_constraint(model, f, MOI.RelativeEntropyCone(2m - 1)) + return AGEBridge{T,F,G,H}(set.k, ceq, relative_entropy_constraint) +end + +function MOI.supports_constraint( + ::Type{<:AGEBridge{T}}, + ::Type{<:MOI.AbstractVectorFunction}, + ::Type{<:AGECone}, +) where {T} + return true +end + +function MOI.Bridges.added_constrained_variable_types(::Type{<:AGEBridge}) + return Tuple{Type}[] +end + +function MOI.Bridges.added_constraint_types(::Type{<:AGEBridge{T,F,G}}) where {T,F,G} + return [(F, MOI.EqualTo{T}), (G, MOI.RelativeEntropyCone)] +end + +function MOI.Bridges.Constraint.concrete_bridge_type( + ::Type{<:AGEBridge{T}}, + H::Type{<:MOI.AbstractVectorFunction}, + ::Type{<:AGECone}, +) where {T} + S = MOI.Utilities.scalar_type(H) + F = MOI.Utilities.promote_operation(+, T, S, MOI.ScalarAffineFunction{Float64}) + G = MOI.Utilities.promote_operation(vcat, T, S, F) + return AGEBridge{T,F,G,H} +end + +function PolyJuMP.bridges( + ::Type{<:MOI.AbstractVectorFunction}, + ::Type{<:AGECone}, +) + return [AGEBridge] +end + +end + +using PolyJuMP +model = Model() +PolyJuMP.setpolymodule!(model, RelativeEntropy) +@constraint(model, motzkin >= 0) + +# As we can see above, even if the constraint is internally stored as a +# `MOI.VectorAffineFunction`-in-`SAGECone`, the constraint is reshaped +# into a `Polynomial`-in-`SOSSet` for printing thanks to the `reshape_set` we +# have implemented. + +model = Model() +@constraint(model, motzkin in RelativeEntropy.AGESet(x^2*y^2)) + +import ECOS +set_optimizer(model, ECOS.Optimizer) +optimize!(model) + +@show termination_status(model)