Refactor SAGE (#100)

* Refactor SAGE * Rename submodule * Rename submodule * Rename submodule * Rename submodule * Rename submodule * Fixes * Fixes * Fix format * Rename
jump-dev · Oct 18, 2023 · 06bc2ea · 06bc2ea
1 parent a5c4802
commit 06bc2ea
Show file tree

Hide file tree

Showing 7 changed files with 103 additions and 141 deletions.
diff --git a/src/PolyJuMP.jl b/src/PolyJuMP.jl
@@ -34,6 +34,6 @@ include("default.jl")
 include("model.jl")
 include("KKT/KKT.jl")
 include("QCQP/QCQP.jl")
-include("RelativeEntropy/RelativeEntropy.jl")
+include("SAGE/SAGE.jl")
 
 end # module
diff --git a/src/RelativeEntropy/RelativeEntropy.jl → src/SAGE/SAGE.jl b/src/RelativeEntropy/RelativeEntropy.jl → src/SAGE/SAGE.jl
@@ -1,47 +1,22 @@
-module RelativeEntropy
+module SAGE
 
 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
+struct Cone{C} <: MOI.AbstractVectorSet
+    cone::C
     α::Matrix{Int}
 end
 
-struct SignomialAGECone <: AbstractAGECone
-    α::Matrix{Int}
-    k::Int
-end
-
-struct PolynomialAGECone <: AbstractAGECone
-    α::Matrix{Int}
-    k::Int
+function JuMP.reshape_set(c::Cone, ::PolyJuMP.PolynomialShape)
+    return c.cone
 end
 
-MOI.dimension(set::AbstractAGECone) = size(set.α, 1)
-Base.copy(set::AbstractAGECone) = set
+MOI.dimension(set::Cone) = size(set.α, 1)
+Base.copy(set::Cone) = set
 
 function _exponents_matrix(monos)
     α = Matrix{Int}(undef, length(monos), MP.nvariables(monos))
@@ -54,64 +29,48 @@ function _exponents_matrix(monos)
     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 Signomials{M<:Union{Nothing,Int,MP.AbstractMonomial}} <: PolyJuMP.PolynomialSet
 
-struct PolynomialSAGESet <: PolyJuMP.PolynomialSet end
-function JuMP.reshape_set(::PolynomialSAGECone, ::PolyJuMP.PolynomialShape)
-    return PolynomialSAGESet()
+**S**ums of **A**M/**G**M **E**xponential for signomials.
+"""
+struct Signomials{M<:Union{Nothing,Int,MP.AbstractMonomial}} <:
+       PolyJuMP.PolynomialSet
+    monomial::M
 end
-function JuMP.moi_set(::PolynomialSAGESet, monos)
-    return PolynomialSAGECone(_exponents_matrix(monos))
+Signomials() = Signomials(nothing)
+_index(_, ::Nothing) = nothing
+_index(monos, mono::MP.AbstractMonomial) = findfirst(isequal(mono), monos)::Int
+function JuMP.moi_set(c::Signomials, monos)
+    return Cone(Signomials(_index(monos, c.monomial)), _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 Polynomials{M<:Union{Nothing,Int,MP.AbstractMonomial}} <: PolyJuMP.PolynomialSet
 
-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)
+**S**ums of **A**M/**G**M **E**xponential for polynomials.
+"""
+struct Polynomials{M<:Union{Nothing,Int,MP.AbstractMonomial}} <:
+       PolyJuMP.PolynomialSet
+    monomial::M
+end
+Polynomials() = Polynomials(nothing)
+function JuMP.moi_set(c::Polynomials, monos)
+    return Cone(
+        Polynomials(_index(monos, c.monomial)),
+        _exponents_matrix(monos),
+    )
 end
 
 function setdefaults!(data::PolyJuMP.Data)
-    PolyJuMP.setdefault!(data, PolyJuMP.NonNegPoly, PolynomialSAGESet)
+    PolyJuMP.setdefault!(data, PolyJuMP.NonNegPoly, Polynomials)
     return
 end
 
 function JuMP.build_constraint(
     _error::Function,
     p,
-    set::Union{
-        SignomialSAGESet,
-        PolynomialSAGESet,
-        SignomialAGESet,
-        PolynomialAGESet,
-    };
+    set::Union{Signomials,Polynomials};
     kws...,
 )
     coefs = PolyJuMP.non_constant_coefficients(p)
@@ -188,7 +147,7 @@ include("bridges/sage.jl")
 
 function PolyJuMP.bridges(
     F::Type{<:MOI.AbstractVectorFunction},
-    ::Type{<:SignomialSAGECone},
+    ::Type{Cone{Signomials{Nothing}}},
 )
     return [(SAGEBridge, PolyJuMP._coef_type(F))]
 end
@@ -197,7 +156,7 @@ include("bridges/age.jl")
 
 function PolyJuMP.bridges(
     F::Type{<:MOI.AbstractVectorFunction},
-    ::Type{<:SignomialAGECone},
+    ::Type{Cone{Signomials{Int}}},
 )
     return [(AGEBridge, PolyJuMP._coef_type(F))]
 end
@@ -206,9 +165,9 @@ include("bridges/signomial.jl")
 
 function PolyJuMP.bridges(
     F::Type{<:MOI.AbstractVectorFunction},
-    ::Type{<:Union{PolynomialSAGECone,PolynomialAGECone}},
-)
-    return [(SignomialBridge, PolyJuMP._coef_type(F))]
+    ::Type{Cone{Polynomials{M}}},
+) where {M}
+    return [(SignomialsBridge, PolyJuMP._coef_type(F))]
 end
 
 end
diff --git a/src/RelativeEntropy/bridges/age.jl → src/SAGE/bridges/age.jl b/src/RelativeEntropy/bridges/age.jl → src/SAGE/bridges/age.jl
@@ -36,7 +36,7 @@ In this bridge, we use the second formulation.
 "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."
+"Signomials and polynomial optimization via relative entropy and partial dualization."
 Mathematical Programming Computation 13 (2021): 257-295.
 https://arxiv.org/pdf/1907.00814.pdf
 """
@@ -50,7 +50,7 @@ function MOI.Bridges.Constraint.bridge_constraint(
     ::Type{AGEBridge{T,F,G,H}},
     model,
     func::H,
-    set::SignomialAGECone,
+    set::Cone{Signomials{Int}},
 ) where {T,F,G,H}
     m = size(set.α, 1)
     ν = MOI.add_variables(model, m - 1)
@@ -59,7 +59,7 @@ function MOI.Bridges.Constraint.bridge_constraint(
         f = zero(typeof(sumν))
         j = 0
         for i in 1:m
-            if i == set.k
+            if i == set.cone.monomial
                 MA.sub_mul!!(f, convert(T, set.α[i, var]), sumν)
             else
                 j += 1
@@ -72,19 +72,23 @@ function MOI.Bridges.Constraint.bridge_constraint(
     f = MOI.Utilities.operate(
         vcat,
         T,
-        scalars[set.k] + sumν,
-        scalars[setdiff(1:m, set.k)],
+        scalars[set.cone.monomial] + sumν,
+        scalars[setdiff(1:m, set.cone.monomial)],
         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)
+    return AGEBridge{T,F,G,H}(
+        set.cone.monomial,
+        ceq,
+        relative_entropy_constraint,
+    )
 end
 
 function MOI.supports_constraint(
     ::Type{<:AGEBridge{T}},
     ::Type{<:MOI.AbstractVectorFunction},
-    ::Type{<:SignomialAGECone},
+    ::Type{Cone{Signomials{Int}}},
 ) where {T}
     return true
 end
@@ -102,7 +106,7 @@ end
 function MOI.Bridges.Constraint.concrete_bridge_type(
     ::Type{<:AGEBridge{T}},
     H::Type{<:MOI.AbstractVectorFunction},
-    ::Type{<:SignomialAGECone},
+    ::Type{Cone{Signomials{Int}}},
 ) where {T}
     S = MOI.Utilities.scalar_type(H)
     F = MOI.Utilities.promote_operation(

diff --git a/src/RelativeEntropy/bridges/sage.jl → src/SAGE/bridges/sage.jl b/src/RelativeEntropy/bridges/sage.jl → src/SAGE/bridges/sage.jl
@@ -1,7 +1,7 @@
 struct SAGEBridge{T,F,G} <: MOI.Bridges.Constraint.AbstractBridge
     ν::Matrix{MOI.VariableIndex}
     age_constraints::Vector{
-        MOI.ConstraintIndex{MOI.VectorOfVariables,SignomialAGECone},
+        MOI.ConstraintIndex{MOI.VectorOfVariables,Cone{Signomials{Int}}},
     }
     equality_constraints::Vector{MOI.ConstraintIndex{F,MOI.EqualTo{T}}}
 end
@@ -10,16 +10,16 @@ function MOI.Bridges.Constraint.bridge_constraint(
     ::Type{SAGEBridge{T,F,G}},
     model,
     func::G,
-    set::SignomialSAGECone,
+    set::Cone{Signomials{Nothing}},
 ) where {T,F,G}
     m = size(set.α, 1)
     ν = Matrix{MOI.VariableIndex}(undef, m, m)
-    A = SignomialAGECone
+    A = Cone{Signomials{Int}}
     age_constraints =
         Vector{MOI.ConstraintIndex{MOI.VectorOfVariables,A}}(undef, m)
     for k in 1:m
         ν[k, :], age_constraints[k] =
-            MOI.add_constrained_variables(model, A(set.α, k))
+            MOI.add_constrained_variables(model, Cone(Signomials(k), set.α))
     end
     scalars = MOI.Utilities.eachscalar(func)
     equality_constraints = map(1:m) do i
@@ -39,13 +39,13 @@ end
 function MOI.supports_constraint(
     ::Type{<:SAGEBridge{T}},
     ::Type{<:MOI.AbstractVectorFunction},
-    ::Type{<:SignomialSAGECone},
+    ::Type{Cone{Signomials{Nothing}}},
 ) where {T}
     return true
 end
 
 function MOI.Bridges.added_constrained_variable_types(::Type{<:SAGEBridge})
-    return Tuple{Type}[(SignomialAGECone,)]
+    return Tuple{Type}[(Cone{Signomials{Int}},)]
 end
 
 function MOI.Bridges.added_constraint_types(
@@ -57,7 +57,7 @@ end
 function MOI.Bridges.Constraint.concrete_bridge_type(
     ::Type{<:SAGEBridge{T}},
     G::Type{<:MOI.AbstractVectorFunction},
-    ::Type{<:SignomialSAGECone},
+    ::Type{Cone{Signomials{Nothing}}},
 ) where {T}
     S = MOI.Utilities.scalar_type(G)
     F = MOI.Utilities.promote_operation(

diff --git a/src/RelativeEntropy/bridges/signomial.jl → src/SAGE/bridges/signomial.jl b/src/RelativeEntropy/bridges/signomial.jl → src/SAGE/bridges/signomial.jl
@@ -1,27 +1,22 @@
 """
-    SignomialBridge{T,S,P,F} <: MOI.Bridges.Constraint.AbstractBridge
+    SignomialsBridge{T,S,P,F} <: MOI.Bridges.Constraint.AbstractBridge
 
-We use the Signomial Representative `SR` equation of [MCW21].
+We use the Signomials Representative `SR` equation of [MCW21].
 
 [MCW20] Riley Murray, Venkat Chandrasekaran, Adam Wierman
 "Newton Polytopes and Relative Entropy Optimization"
 https://arxiv.org/abs/1810.01614
 [MCW21] Murray, Riley, Venkat Chandrasekaran, and Adam Wierman.
-"Signomial and polynomial optimization via relative entropy and partial dualization."
+"Signomials and polynomial optimization via relative entropy and partial dualization."
 Mathematical Programming Computation 13 (2021): 257-295.
 https://arxiv.org/pdf/1907.00814.pdf
 """
-struct SignomialBridge{T,S,P,F} <: MOI.Bridges.Constraint.AbstractBridge
+struct SignomialsBridge{T,S,P,F} <: MOI.Bridges.Constraint.AbstractBridge
     constraint::MOI.ConstraintIndex{F,S}
 end
 
-_signomial(set::PolynomialSAGECone) = SignomialSAGECone(set.α)
-_signomial(::Type{PolynomialSAGECone}) = SignomialSAGECone
-_signomial(set::PolynomialAGECone) = SignomialAGECone(set.α, set.k)
-_signomial(::Type{PolynomialAGECone}) = SignomialAGECone
-
 function MOI.Bridges.Constraint.bridge_constraint(
-    ::Type{SignomialBridge{T,S,P,F}},
+    ::Type{SignomialsBridge{T,S,P,F}},
     model,
     func::F,
     set,
@@ -44,42 +39,46 @@ function MOI.Bridges.Constraint.bridge_constraint(
             g[i] = vi
         end
     end
-    constraint =
-        MOI.add_constraint(model, MOI.Utilities.vectorize(g), _signomial(set))
-    return SignomialBridge{T,S,P,F}(constraint)
+    constraint = MOI.add_constraint(
+        model,
+        MOI.Utilities.vectorize(g),
+        Cone(Signomials(set.cone.monomial), set.α),
+    )
+    return SignomialsBridge{T,S,P,F}(constraint)
 end
 
 function MOI.supports_constraint(
-    ::Type{<:SignomialBridge{T}},
+    ::Type{<:SignomialsBridge{T}},
     ::Type{<:MOI.AbstractVectorFunction},
-    ::Type{<:Union{PolynomialSAGECone,PolynomialAGECone}},
-) where {T}
+    ::Type{Cone{Polynomials{M}}},
+) where {T,M}
     return true
 end
 
-function MOI.Bridges.added_constrained_variable_types(::Type{<:SignomialBridge})
+function MOI.Bridges.added_constrained_variable_types(
+    ::Type{<:SignomialsBridge},
+)
     return Tuple{Type}[(MOI.Reals,)]
 end
 
 function MOI.Bridges.added_constraint_types(
-    ::Type{<:SignomialBridge{T,S,P,F}},
+    ::Type{<:SignomialsBridge{T,S,P,F}},
 ) where {T,S,P,F}
     return [(F, S)]
 end
 
 function MOI.Bridges.Constraint.concrete_bridge_type(
-    ::Type{<:SignomialBridge{T}},
+    ::Type{<:SignomialsBridge{T}},
     F::Type{<:MOI.AbstractVectorFunction},
-    P::Type{<:Union{PolynomialSAGECone,PolynomialAGECone}},
-) where {T}
-    S = _signomial(P)
-    return SignomialBridge{T,S,P,F}
+    P::Type{Cone{Polynomials{M}}},
+) where {T,M}
+    return SignomialsBridge{T,Cone{Signomials{M}},P,F}
 end
 
 function MOI.get(
     model::MOI.ModelLike,
     attr::DecompositionAttribute,
-    bridge::SignomialBridge,
+    bridge::SignomialsBridge,
 )
     return MOI.get(model, attr, bridge.constraint)
 end