Add RelativeEntropy/SAGE solver (#99)

* Add RelativeEntropy/SAGE solver

* Fix format

* Fix

* Add polynomial SAGE

* Fix format

* Add tests

* Add decomposition and fix test

* Fix format

src/PolyJuMP.jl

@@ -34,5 +34,6 @@ include("default.jl")
 include("model.jl")
 include("KKT/KKT.jl")
 include("QCQP/QCQP.jl")
+include("RelativeEntropy/RelativeEntropy.jl")
 
 end # module

src/RelativeEntropy/RelativeEntropy.jl

@@ -0,0 +1,214 @@
+module RelativeEntropy
+
+import MutableArithmetics as MA
+import MultivariatePolynomials as MP
+import MathOptInterface as MOI
+import JuMP
+import PolyJuMP
+
+abstract type AbstractAGECone <: MOI.AbstractVectorSet end
+
+"""
+    struct SignomialSAGECone <: MOI.AbstractVectorSet
+        α::Matrix{Int}
+    end
+
+**S**ums of **A**M/**G**M **E**xponential for signomials.
+"""
+struct SignomialSAGECone <: AbstractAGECone
+    α::Matrix{Int}
+end
+
+"""
+    struct PolynomialSAGECone <: MOI.AbstractVectorSet
+        α::Matrix{Int}
+    end
+
+**S**ums of **A**M/**G**M **E**xponential for polynomials.
+"""
+struct PolynomialSAGECone <: AbstractAGECone
+    α::Matrix{Int}
+end
+
+struct SignomialAGECone <: AbstractAGECone
+    α::Matrix{Int}
+    k::Int
+end
+
+struct PolynomialAGECone <: AbstractAGECone
+    α::Matrix{Int}
+    k::Int
+end
+
+MOI.dimension(set::AbstractAGECone) = size(set.α, 1)
+Base.copy(set::AbstractAGECone) = set
+
+function _exponents_matrix(monos)
+    α = Matrix{Int}(undef, length(monos), MP.nvariables(monos))
+    for (i, mono) in enumerate(monos)
+        exp = MP.exponents(mono)
+        for j in eachindex(exp)
+            α[i, j] = exp[j]
+        end
+    end
+    return α
+end
+
+struct SignomialSAGESet <: PolyJuMP.PolynomialSet end
+function JuMP.reshape_set(::SignomialSAGECone, ::PolyJuMP.PolynomialShape)
+    return SignomialSAGESet()
+end
+function JuMP.moi_set(::SignomialSAGESet, monos)
+    return SignomialSAGECone(_exponents_matrix(monos))
+end
+
+struct PolynomialSAGESet <: PolyJuMP.PolynomialSet end
+function JuMP.reshape_set(::PolynomialSAGECone, ::PolyJuMP.PolynomialShape)
+    return PolynomialSAGESet()
+end
+function JuMP.moi_set(::PolynomialSAGESet, monos)
+    return PolynomialSAGECone(_exponents_matrix(monos))
+end
+
+struct SignomialAGESet{MT<:MP.AbstractMonomial} <: PolyJuMP.PolynomialSet
+    monomial::MT
+end
+function JuMP.reshape_set(
+    set::SignomialAGECone,
+    shape::PolyJuMP.PolynomialShape,
+)
+    return SignomialAGESet(shape.monomials[set.k])
+end
+function JuMP.moi_set(set::SignomialAGESet, monos)
+    k = findfirst(isequal(set.monomial), monos)
+    return SignomialAGECone(_exponents_matrix(monos), k)
+end
+
+struct PolynomialAGESet{MT<:MP.AbstractMonomial} <: PolyJuMP.PolynomialSet
+    monomial::MT
+end
+function JuMP.reshape_set(
+    set::PolynomialAGECone,
+    shape::PolyJuMP.PolynomialShape,
+)
+    return PolynomialAGESet(shape.monomials[set.k])
+end
+function JuMP.moi_set(set::PolynomialAGESet, monos)
+    k = findfirst(isequal(set.monomial), monos)
+    return PolynomialAGECone(_exponents_matrix(monos), k)
+end
+
+function setdefaults!(data::PolyJuMP.Data)
+    PolyJuMP.setdefault!(data, PolyJuMP.NonNegPoly, PolynomialSAGESet)
+    return
+end
+
+function JuMP.build_constraint(
+    _error::Function,
+    p,
+    set::Union{
+        SignomialSAGESet,
+        PolynomialSAGESet,
+        SignomialAGESet,
+        PolynomialAGESet,
+    };
+    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
+
+const APL{T} = MP.AbstractPolynomialLike{T}
+
+"""
+    struct Decomposition{T, PT}
+
+Represents a SAGE decomposition.
+"""
+struct Decomposition{T,P<:APL{T}} <: APL{T}
+    polynomials::Vector{P}
+    function Decomposition(ps::Vector{P}) where {T,P<:APL{T}}
+        return new{T,P}(ps)
+    end
+end
+
+function Base.show(io::IO, p::Decomposition)
+    for (i, q) in enumerate(p.polynomials)
+        print(io, "(")
+        print(io, q)
+        print(io, ")")
+        if i != length(p.polynomials)
+            print(io, " + ")
+        end
+    end
+end
+
+function MP.polynomial(d::Decomposition)
+    return sum(d.polynomials)
+end
+
+"""
+    struct DecompositionAttribute{T} <: MOI.AbstractConstraintAttribute
+        tol::T
+    end
+
+A constraint attribute for the [`Decomposition`](@ref) of a constraint.
+"""
+struct DecompositionAttribute{T} <: MOI.AbstractConstraintAttribute
+    tol::T
+    result_index::Int
+end
+function DecompositionAttribute(tol::Real)
+    return DecompositionAttribute(tol, 1)
+end
+
+function decomposition(
+    con_ref::JuMP.ConstraintRef;
+    tol::Real,
+    result_index::Int = 1,
+)
+    monos = con_ref.shape.monomials
+    attr = DecompositionAttribute(tol, result_index)
+    return Decomposition([
+        MP.polynomial(a, monos) for
+        a in MOI.get(JuMP.owner_model(con_ref), attr, con_ref)
+    ])
+end
+
+MOI.is_set_by_optimize(::DecompositionAttribute) = true
+
+include("bridges/sage.jl")
+
+function PolyJuMP.bridges(
+    F::Type{<:MOI.AbstractVectorFunction},
+    ::Type{<:SignomialSAGECone},
+)
+    return [(SAGEBridge, PolyJuMP._coef_type(F))]
+end
+
+include("bridges/age.jl")
+
+function PolyJuMP.bridges(
+    F::Type{<:MOI.AbstractVectorFunction},
+    ::Type{<:SignomialAGECone},
+)
+    return [(AGEBridge, PolyJuMP._coef_type(F))]
+end
+
+include("bridges/signomial.jl")
+
+function PolyJuMP.bridges(
+    F::Type{<:MOI.AbstractVectorFunction},
+    ::Type{<:Union{PolynomialSAGECone,PolynomialAGECone}},
+)
+    return [(SignomialBridge, PolyJuMP._coef_type(F))]
+end
+
+end

src/RelativeEntropy/bridges/age.jl

@@ -0,0 +1,116 @@
+"""
+    AGEBridge{T,F,G,H} <: MOI.Bridges.Constraint.AbstractBridge
+
+The nonnegativity of `x ≥ 0` in
+```
+∑ ci x^αi ≥ -c0 x^α0
+```
+can be reformulated as
+```
+∑ ci exp(αi'y) ≥ -β exp(α0'y)
+```
+In any case, it is shown to be equivalent to
+```
+∃ ν ≥ 0 s.t. D(ν, exp(1)*c) ≤ β, ∑ νi αi = α0 ∑ νi [CP16, (3)]
+```
+where `N(ν, λ) = ∑ νj log(νj/λj)` is the relative entropy function.
+The constant `exp(1)` can also be moved out of `D` into
+```
+∃ ν ≥ 0 s.t. D(ν, c) - ∑ νi ≤ β, ∑ νi αi = α0 ∑ νi [MCW21, (2)]
+```
+The relative entropy cone in MOI is `(u, v, w)` such that `D(w, v) ≥ u`.
+Therefore, we can either encode `(β, exp(1)*c, ν)` [CP16, (3)] or
+`(β + ∑ νi, c, ν)` [MCW21, (2)].
+In this bridge, we use the second formulation.
+
+!!! note
+    A direct application of the Arithmetic-Geometric mean inequality gives
+    ```
+    ∃ ν ≥ 0 s.t. D(ν, exp(1)*c) ≤ -log(-β), ∑ νi αi = α0, ∑ νi = 1 [CP16, (4)]
+    ```
+    which is not jointly convex in (ν, c, β).
+    The key to get the convex formulation [CP16, (3)] or [MCW21, (2)] instead is to
+    use the convex conjugacy between the exponential and the negative entropy functions [CP16, (iv)].
+
+[CP16] Chandrasekaran, Venkat, and Parikshit Shah.
+"Relative entropy relaxations for signomial optimization."
+SIAM Journal on Optimization 26.2 (2016): 1147-1173.
+[MCW21] Murray, Riley, Venkat Chandrasekaran, and Adam Wierman.
+"Signomial and polynomial optimization via relative entropy and partial dualization."
+Mathematical Programming Computation 13 (2021): 257-295.
+https://arxiv.org/pdf/1907.00814.pdf
+"""
+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
+
+function MOI.Bridges.Constraint.bridge_constraint(
+    ::Type{AGEBridge{T,F,G,H}},
+    model,
+    func::H,
+    set::SignomialAGECone,
+) 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 = zero(typeof(sumν))
+        j = 0
+        for i in 1:m
+            if i == set.k
+                MA.sub_mul!!(f, convert(T, set.α[i, var]), sumν)
+            else
+                j += 1
+                MA.add_mul!!(f, convert(T, set.α[i, var]), ν[j])
+            end
+        end
+        return MOI.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{<:SignomialAGECone},
+) 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{<:SignomialAGECone},
+) 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