[breaking] refactor RelativeEntropyEpiCone into RelativeEntropyEpiConeSquare (#663)

* [breaking] refactor RelativeEntropyEpiCone into
RelativeEntropyEpiConeSquare

* Update

* Update

* Update

* Update

* Update

Author: odow

docs/src/manual/operations.md

@@ -123,7 +123,7 @@ solver (including SCS and Mosek).
 | `trace_logm(X, C)`                | $\textrm{Tr}(C \log X)$            | concave in X | not monotonic | IC: $X \succeq 0$, $C \succeq 0$, $C$ constant; uses natural log |
 | `trace_mpower(A, t, C)`           | $\textrm{Tr}(C A^t)$ | concave in A for $t \in [0,1]$, convex for $t \in [-1,0] \cup [1,2]$   | not monotonic | IC: $X \succeq 0$, $C \succeq 0$, $C$ constant, $t \in [-1, 2]$ |
 | `lieb_ando(A, B, K, t)`           | $\textrm{Tr}(K' A^{1-t} K B^t)$    | concave in A,B for $t \in [0,1]$, convex for $t \in [-1,0] \cup [1,2]$ | not monotonic | IC: $A \succeq 0$, $B \succeq 0$, $K$ constant, $t \in [-1, 2]$ |
-| `T in RelativeEntropyEpiCone(X, Y, m, k, e)` | $T \succeq e' X^{1/2} \log(X^{1/2} Y^{-1} X^{1/2}) X^{1/2} e$ | convex | not monotonic | IC: $e$ constant; uses natural log |
+| `Convex.GenericConstraint((T, X, Y), RelativeEntropyEpiConeSquare(size(X, 1), m, k, e))` | $T \succeq e' X^{1/2} \log(X^{1/2} Y^{-1} X^{1/2}) X^{1/2} e$ | convex | not monotonic | IC: $e$ constant; uses natural log |
 
 ## Exponential + SDP representable Functions
 

src/Convex.jl

@@ -46,7 +46,7 @@ export conv,
     sumsquares,
     GeometricMeanHypoConeSquare,
     GeometricMeanEpiConeSquare,
-    RelativeEntropyEpiCone,
+    RelativeEntropyEpiConeSquare,
     quantum_relative_entropy,
     quantum_entropy,
     trace_logm,

src/atoms/QuantumEntropyAtom.jl

@@ -81,7 +81,8 @@ function new_conic_form!(
         Variable(size(X))
     end
     I = Matrix(one(T) * LinearAlgebra.I(size(X, 1)))
-    add_constraint!(context, τ in RelativeEntropyEpiCone(X, I, atom.m, atom.k))
+    set = RelativeEntropyEpiConeSquare(size(X, 1), atom.m, atom.k)
+    add_constraint!(context, GenericConstraint((τ, X, I), set))
     # It's already a real mathematically, but need to make it a real type.
     return conic_form!(context, real(-LinearAlgebra.tr(τ)))
 end

src/atoms/QuantumRelativeEntropyAtom.jl

@@ -93,12 +93,9 @@ function new_conic_form!(
     add_constraint!(context, A ⪰ 0)
     add_constraint!(context, B ⪰ 0)
     I = Matrix(one(T) * LinearAlgebra.I(size(A, 1)))
-    m, k, e = atom.m, atom.k, vec(I)
-    τ = Variable()
-    add_constraint!(
-        context,
-        τ in RelativeEntropyEpiCone(kron(A, I), kron(I, conj(B)), m, k, e),
-    )
+    τ, X, Y = Variable(), kron(A, I), kron(I, conj(B))
+    set = RelativeEntropyEpiConeSquare(size(X, 1), atom.m, atom.k, vec(I))
+    add_constraint!(context, GenericConstraint((τ, X, Y), set))
     return conic_form!(context, τ)
 end
 

src/atoms/TraceLogmAtom.jl

@@ -102,7 +102,8 @@ function new_conic_form!(context::Context{T}, atom::TraceLogmAtom) where {T}
         Variable(size(X))
     end
     I = Matrix(one(T) * LinearAlgebra.I(size(X, 1)))
-    add_constraint!(context, τ in RelativeEntropyEpiCone(I, X, atom.m, atom.k))
+    set = RelativeEntropyEpiConeSquare(size(X, 1), atom.m, atom.k)
+    add_constraint!(context, GenericConstraint((τ, I, X), set))
     # It's already a real mathematically, but need to make it a real type.
     return conic_form!(context, real(-LinearAlgebra.tr(atom.C * τ)))
 end

src/constraints/RelativeEntropyEpiCone.jl

@@ -5,154 +5,110 @@
 # in the LICENSE.md file or at https://opensource.org/licenses/MIT.
 
 """
-Constrains τ to
-  τ ⪰ e' * X^{1/2} * logm(X^{1/2}*Y^{-1}*X^{1/2}) * X^{1/2} * e
-
-This function implements the semidefinite programming approximation given in
-the reference below.  Parameters m and k control the accuracy of this
-approximation: m is the number of quadrature nodes to use and k the number
-of square-roots to take. See reference for more details.
-
-All expressions and atoms are subtypes of AbstractExpr.
-Please read expressions.jl first.
-
-REFERENCE
-  Ported from CVXQUAD which is based on the paper: "Semidefinite
-  approximations of matrix logarithm" by Hamza Fawzi, James Saunderson and
-  Pablo A. Parrilo (arXiv:1705.00812)
-"""
-mutable struct RelativeEntropyEpiCone
-    X::AbstractExpr
-    Y::AbstractExpr
-    m::Integer
-    k::Integer
-    e::AbstractMatrix
-    size::Tuple{Int,Int}
-
-    function RelativeEntropyEpiCone(
-        X::AbstractExpr,
-        Y::AbstractExpr,
+    RelativeEntropyEpiConeSquare(
+        side_dimension::Int,
         m::Integer = 3,
         k::Integer = 3,
-        e::AbstractArray = Matrix(1.0 * LinearAlgebra.I, size(X)),
+        e::AbstractArray = Matrix(1.0 * LinearAlgebra.I(side_dimension)),
     )
-        if size(X) != size(Y)
-            throw(DimensionMismatch("X and Y must be the same size"))
-        end
-        n = size(X)[1]
-        if size(X) != (n, n)
-            throw(DimensionMismatch("X and Y must be square"))
-        end
-        if length(size(e)) == 1
-            e = reshape(e, (size(e)[1], 1))
-        end
-        erows, ecols = size(e)
-        if ndims(e) != 2 || erows != n
-            throw(DimensionMismatch("e matrix must have n rows"))
-        end
-        return new(X, Y, m, k, e, (ecols, ecols))
-    end
 
-    function RelativeEntropyEpiCone(
-        X::Value,
-        Y::AbstractExpr,
-        m::Integer = 3,
-        k::Integer = 3,
-        e::AbstractArray = Matrix(1.0 * LinearAlgebra.I, size(X)),
-    )
-        return RelativeEntropyEpiCone(constant(X), Y, m, k, e)
-    end
+Constrains `(τ, X, Y)` to:
+```
+τ ⪰ e' * X^{1/2} * logm(X^{1/2}*Y^{-1}*X^{1/2}) * X^{1/2} * e
+```
 
-    function RelativeEntropyEpiCone(
-        X::AbstractExpr,
-        Y::Value,
-        m::Integer = 3,
-        k::Integer = 3,
-        e::AbstractArray = Matrix(1.0 * LinearAlgebra.I, size(X)),
-    )
-        return RelativeEntropyEpiCone(X, constant(Y), m, k, e)
-    end
+This set implements the semidefinite programming approximation given in the
+reference below.
 
-    function RelativeEntropyEpiCone(
-        X::Value,
-        Y::Value,
-        m::Integer = 3,
-        k::Integer = 3,
-        e::AbstractArray = Matrix(1.0 * LinearAlgebra.I, size(X)),
-    )
-        return RelativeEntropyEpiCone(constant(X), constant(Y), m, k, e)
-    end
-end
+Parameters `m` and `k` control the accuracy of this approximation: `m` is the
+number of quadrature nodes to use and `k` the number of square-roots to take.
+See reference for more details.
 
-mutable struct RelativeEntropyEpiConeConstraint <: Constraint
-    τ::AbstractExpr
-    cone::RelativeEntropyEpiCone
+## Reference
 
-    function RelativeEntropyEpiConeConstraint(
-        τ::AbstractExpr,
-        cone::RelativeEntropyEpiCone,
+Ported from CVXQUAD which is based on the paper: "Semidefinite approximations of
+matrix logarithm" by Hamza Fawzi, James Saunderson and Pablo A. Parrilo
+(arXiv:1705.00812)
+"""
+struct RelativeEntropyEpiConeSquare{T} <: MOI.AbstractVectorSet
+    side_dimension::Int
+    m::Int
+    k::Int
+    e::Matrix{T}
+
+    function RelativeEntropyEpiConeSquare(
+        side_dimension::Int,
+        m::Integer = 3,
+        k::Integer = 3,
+        e::AbstractArray = Matrix(1.0 * LinearAlgebra.I(side_dimension)),
     )
-        if size(τ) != cone.size
-            throw(DimensionMismatch("τ must be size $(cone.size)"))
+        if length(size(e)) == 1
+            e = reshape(e, (size(e, 1), 1))
         end
-        return new(τ, cone)
-    end
-
-    function RelativeEntropyEpiConeConstraint(
-        τ::Value,
-        cone::RelativeEntropyEpiCone,
-    )
-        return RelativeEntropyEpiConeConstraint(constant(τ), cone)
+        if ndims(e) != 2 || size(e, 1) != side_dimension
+            throw(DimensionMismatch("e matrix must have $side_dimension rows"))
+        end
+        return new{eltype(e)}(side_dimension, m, k, e)
     end
 end
 
-function iscomplex(c::RelativeEntropyEpiConeConstraint)
-    return iscomplex(c.τ) || iscomplex(c.cone)
+function MOI.dimension(set::RelativeEntropyEpiConeSquare)
+    return size(set.e, 2)^2 + 2 * set.side_dimension^2
 end
 
-iscomplex(c::RelativeEntropyEpiCone) = iscomplex(c.X) || iscomplex(c.Y)
-
-function head(io::IO, ::RelativeEntropyEpiConeConstraint)
-    return print(io, "∈(RelativeEntropyEpiCone)")
+function head(io::IO, ::RelativeEntropyEpiConeSquare)
+    return print(io, "RelativeEntropyEpiConeSquare")
 end
 
-function Base.in(τ, cone::RelativeEntropyEpiCone)
-    return RelativeEntropyEpiConeConstraint(τ, cone)
+function GenericConstraint(func::Tuple, set::RelativeEntropyEpiConeSquare)
+    @assert length(func) == 3
+    sizes = (size(set.e, 2), set.side_dimension, set.side_dimension)
+    for (i, f) in enumerate(func)
+        n = LinearAlgebra.checksquare(f)
+        if n != sizes[i]
+            throw(
+                DimensionMismatch(
+                    "Matrix of side dimension `$n` does not match required side dimension `$(sizes[i])`",
+                ),
+            )
+        end
+    end
+    return GenericConstraint(vcat(vec.(func)...), set)
 end
 
-function AbstractTrees.children(constraint::RelativeEntropyEpiConeConstraint)
-    return (constraint.τ, constraint.cone.X, constraint.cone.Y)
+function _get_matrices(c::GenericConstraint{<:RelativeEntropyEpiConeSquare})
+    n_τ, n_x = size(c.set.e, 2), c.set.side_dimension
+    d_τ, d_x = n_τ^2, n_x^2
+    τ = reshape(c.child[1:d_τ], n_τ, n_τ)
+    X = reshape(c.child[d_τ.+(1:d_x)], n_x, n_x)
+    Y = reshape(c.child[d_τ.+d_x.+(1:d_x)], n_x, n_x)
+    return τ, X, Y
 end
 
 # This negative relative entropy function is matrix convex (arxiv:1705.00812).
 # So if X and Y are convex sets, then τ ⪰ -D_op(X || Y) will be a convex set.
-function vexity(constraint::RelativeEntropyEpiConeConstraint)
-    X = vexity(constraint.cone.X)
-    Y = vexity(constraint.cone.Y)
-    τ = vexity(constraint.τ)
-    # NOTE: can't say X == NotDcp() because the NotDcp constructor prints a warning message.
-    if typeof(X) == ConcaveVexity || typeof(X) == NotDcp
-        return NotDcp()
-    end
-    if typeof(Y) == ConcaveVexity || typeof(Y) == NotDcp
-        return NotDcp()
-    end
-    vex = ConvexVexity() + (-τ)
-    if vex == ConcaveVexity()
+function vexity(constraint::GenericConstraint{<:RelativeEntropyEpiConeSquare})
+    τ, X, Y = _get_matrices(constraint)
+    if vexity(X) in (ConcaveVexity(), NotDcp()) ||
+       vexity(Y) in (ConcaveVexity(), NotDcp())
         return NotDcp()
     end
-    return vex
+    return ConvexVexity() + -vexity(τ)
 end
 
-function glquad(m)
-    # Compute Gauss-Legendre quadrature nodes and weights on [0, 1]
-    # Code below is from [Trefethen, "Is Gauss quadrature better than
-    # Clenshaw-Curtis?", SIAM Review 2008] and computes the weights and
-    # nodes on [-1, 1].
-    beta = [0.5 ./ sqrt(1 - (2 * i)^-2) for i in 1:(m-1)] # 3-term recurrence coeffs
-    T = LinearAlgebra.diagm(1 => beta, -1 => beta) # Jacobi matrix
+"""
+Compute Gauss-Legendre quadrature nodes and weights on [0, 1].
+
+Code below is from Trefethen (2008), "Is Gauss quadrature better than
+Clenshaw-Curtis?", SIAM Review, and computes the weights and nodes on [-1, 1].
+"""
+function _gauss_legendre_quadrature(m)
+    # 3-term recurrence coeffs
+    beta = [0.5 ./ sqrt(1 - (2 * i)^-2) for i in 1:(m-1)]
+    # Jacobi matrix
+    T = LinearAlgebra.diagm(1 => beta, -1 => beta)
     s, V = LinearAlgebra.eigen(T)
-    w = 2 * V[1, :] .^ 2 # weights
+    w = 2 * V[1, :] .^ 2
     # Translate and scale to [0, 1]
     s = (s .+ 1) / 2
     w = w' / 2
@@ -161,17 +117,12 @@ end
 
 function _add_constraint!(
     context::Context,
-    constraint::RelativeEntropyEpiConeConstraint,
+    constraint::GenericConstraint{<:RelativeEntropyEpiConeSquare},
 )
-    X = constraint.cone.X
-    Y = constraint.cone.Y
-    m = constraint.cone.m
-    k = constraint.cone.k
-    e = constraint.cone.e
-    τ = constraint.τ
-    n = size(X)[1]
-    r = size(e)[2]
-    s, w = glquad(m)
+    τ, X, Y = _get_matrices(constraint)
+    m, k, e = constraint.set.m, constraint.set.k, constraint.set.e
+    n, r = size(X, 1), size(e, 2)
+    s, w = _gauss_legendre_quadrature(m)
     is_complex =
         sign(X) == ComplexSign() ||
         sign(Y) == ComplexSign() ||
@@ -183,25 +134,17 @@ function _add_constraint!(
         Z = Variable(n, n)
         T = [Variable(r, r) for i in 1:m]
     end
-    add_constraint!(
-        context,
-        GenericConstraint(
-            (Z, X, Y),
-            GeometricMeanHypoConeSquare(1 // (2^k), n, false),
-        ),
-    )
+    set = GeometricMeanHypoConeSquare(1 // (2^k), n, false)
+    add_constraint!(context, GenericConstraint((Z, X, Y), set))
     for ii in 1:m
-        # Note that we are dividing by w here because it is easier
-        # to do this than to do sum w_i T(:,...,:,ii) later (cf. line that
-        # involves τ)
-        add_constraint!(
-            context,
-            [
-                e'*X*e-s[ii]*T[ii]/w[ii] e'*X
-                X*e (1-s[ii])*X+s[ii]*Z
-            ] ⪰ 0,
-        )
+        # Note that we are dividing by w here because it is easier to do this
+        # than to do sum w_i T(:,...,:,ii) later (cf. line that involves τ)
+        A_ii = [
+            e'*X*e-s[ii]*T[ii]/w[ii] e'*X
+            X*e (1-s[ii])*X+s[ii]*Z
+        ]
+        add_constraint!(context, A_ii ⪰ 0)
     end
-    add_constraint!(context, (2^k) * sum(T) + τ ⪰ 0)
+    add_constraint!(context, 2^k * sum(T) + τ ⪰ 0)
     return
 end