Add decomposition and fix test

src/RelativeEntropy/RelativeEntropy.jl

@@ -125,6 +125,61 @@ function JuMP.build_constraint(
     )
 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(

src/RelativeEntropy/bridges/sage.jl

@@ -15,13 +15,14 @@ function MOI.Bridges.Constraint.bridge_constraint(
     m = size(set.α, 1)
     ν = Matrix{MOI.VariableIndex}(undef, m, m)
     A = SignomialAGECone
-    c = Vector{MOI.ConstraintIndex{MOI.VectorOfVariables,A}}(undef, m)
+    age_constraints =
+        Vector{MOI.ConstraintIndex{MOI.VectorOfVariables,A}}(undef, m)
     for k in 1:m
-        ν[k, :], c[k] = MOI.add_constrained_variables(model, A(set.α, k))
+        ν[k, :], age_constraints[k] =
+            MOI.add_constrained_variables(model, A(set.α, k))
     end
     scalars = MOI.Utilities.eachscalar(func)
-    n = size(set.α, 2)
-    equality_constraints = map(1:n) do i
+    equality_constraints = map(1:m) do i
         f = MOI.ScalarAffineFunction(
             MOI.ScalarAffineTerm.(one(T), ν[:, i]),
             zero(T),
@@ -32,7 +33,7 @@ function MOI.Bridges.Constraint.bridge_constraint(
             MOI.EqualTo(zero(T)),
         )
     end
-    return SAGEBridge{T,F,G}(ν, c, equality_constraints)
+    return SAGEBridge{T,F,G}(ν, age_constraints, equality_constraints)
 end
 
 function MOI.supports_constraint(
@@ -67,3 +68,20 @@ function MOI.Bridges.Constraint.concrete_bridge_type(
     )
     return SAGEBridge{T,F,G}
 end
+
+function MOI.get(
+    model::MOI.ModelLike,
+    attr::DecompositionAttribute,
+    bridge::SAGEBridge,
+)
+    return filter!(
+        !isempty,
+        [
+            filter(
+                x -> !isapprox(x, zero(x), atol=attr.tol),
+                MOI.get(model, MOI.VariablePrimal(attr.result_index), bridge.ν[k, :])
+            )
+            for k in axes(bridge.ν, 1)
+        ]
+    )
+end

src/RelativeEntropy/bridges/signomial.jl

@@ -12,7 +12,7 @@ Mathematical Programming Computation 13 (2021): 257-295.
 https://arxiv.org/pdf/1907.00814.pdf
 """
 struct SignomialBridge{T,S,P,F} <: MOI.Bridges.Constraint.AbstractBridge
-    constraints::MOI.ConstraintIndex{F,S}
+    constraint::MOI.ConstraintIndex{F,S}
 end
 
 _signomial(set::PolynomialSAGECone) = SignomialSAGECone(set.α)
@@ -36,7 +36,6 @@ function MOI.Bridges.Constraint.bridge_constraint(
             g[i] = vi
         end
     end
-    @show g
     constraint = MOI.add_constraint(model, MOI.Utilities.vectorize(g), _signomial(set))
     return SignomialBridge{T,S,P,F}(constraint)
 end
@@ -67,3 +66,11 @@ function MOI.Bridges.Constraint.concrete_bridge_type(
     S = _signomial(P)
     return SignomialBridge{T,S,P,F}
 end
+
+function MOI.get(
+    model::MOI.ModelLike,
+    attr::DecompositionAttribute,
+    bridge::SignomialBridge,
+)
+    return MOI.get(model, attr, bridge.constraint)
+end

test/relative_entropy.jl

@@ -21,16 +21,23 @@ function _test_motzkin(x, y, T, solver, set, feasible, square, neg)
     else
         motzkin = a^2 * b + a * b^2 + one(T) - 3a * b
     end
-    @show motzkin
-    @constraint(model, motzkin in set)
+    con_ref = @constraint(model, motzkin in set)
     optimize!(model)
-    inner = model.moi_backend.optimizer.model
-    println(inner)
-    vis = MOI.get(inner, MOI.ListOfVariableIndices())
-    @show MOI.get(inner, MOI.VariablePrimal(), vis)
     if feasible
         @test termination_status(model) == MOI.OPTIMAL
         @test primal_status(model) == MOI.FEASIBLE_POINT
+        if set isa Union{PolyJuMP.RelativeEntropy.SignomialSAGESet,
+            PolyJuMP.RelativeEntropy.PolynomialSAGESet}
+            d = PolyJuMP.RelativeEntropy.decomposition(con_ref; tol = 1e-6)
+            p = MP.polynomial(d)
+            if set isa PolyJuMP.RelativeEntropy.SignomialSAGESet
+                @test p ≈ motzkin atol = 1e-6
+            else
+                for m in MP.monomials(p - motzkin)
+                    @test MP.coefficient(p, m) ≈ MP.coefficient(motzkin, m) atol = 1e-6
+                end
+            end
+        end
     else
         @test termination_status(model) == MOI.INFEASIBLE
     end