Refactor GeomMeanEpiCone into GenericConstraint (#604)

* Refactor GeomMeanEpiCone into GenericConstraint

* Fixes

* Fix tests

* Update

* Update

* eye -> I

* Revert Tuple approach

* Fix format

* Update constant.jl

* Update

* Remove sdp_geom_mean_epicone_argcheck

* Update test/test_constraints.jl

Co-authored-by: Benoît Legat <benoit.legat@gmail.com>

* Update

---------

Co-authored-by: odow <o.dowson@gmail.com>
Co-authored-by: Oscar Dowson <odow@users.noreply.github.com>

docs/src/manual/operations.md

@@ -117,7 +117,7 @@ solver (including SCS and Mosek).
 | `matrixfrac(x, P)`                | $x^TP^{-1}x$                       | convex       | not monotonic | IC: P is positive semidefinite |
 | `sumlargesteigs(x, k)`            | sum of top $k$ eigenvalues of $x$  | convex       | not monotonic | IC: P symmetric        |
 | `T in GeomMeanHypoCone(A, B, t)`  | $T \preceq A \#_t B = A^{1/2} (A^{-1/2} B A^{-1/2})^t A^{1/2}$ | concave | increasing    | IC: $A \succeq 0$, $B \succeq 0$, $t \in [0,1]$ |
-| `T in GeomMeanEpiCone(A, B, t)`   | $T \succeq A \#_t B = A^{1/2} (A^{-1/2} B A^{-1/2})^t A^{1/2}$ | convex  | not monotonic | IC: $A \succeq 0$, $B \succeq 0$, $t \in [-1, 0] \cup [1, 2]$ |
+| `Convex.GenericConstraint((T, A, B), GeometricMeanEpiConeSquare(t, size(T, 1)))`   | $T \succeq A \#_t B = A^{1/2} (A^{-1/2} B A^{-1/2})^t A^{1/2}$ | convex  | not monotonic | IC: $A \succeq 0$, $B \succeq 0$, $t \in [-1, 0] \cup [1, 2]$ |
 | `quantum_entropy(X)`              | $-\textrm{Tr}(X \log X)$           | concave      | not monotonic | IC: $X \succeq 0$; uses natural log |
 | `quantum_relative_entropy(A, B)`  | $\textrm{Tr}(A \log A - A \log B)$ | convex       | not monotonic | IC: $A \succeq 0$, $B \succeq 0$; uses natural log |
 | `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 |

src/Convex.jl

@@ -45,7 +45,7 @@ export conv,
     sumsmallest,
     sumsquares,
     GeomMeanHypoCone,
-    GeomMeanEpiCone,
+    GeometricMeanEpiConeSquare,
     RelativeEntropyEpiCone,
     quantum_relative_entropy,
     quantum_entropy,

src/atoms/TraceMpowerAtom.jl

@@ -95,7 +95,13 @@ function new_conic_form!(context::Context{T}, atom::TraceMpowerAtom) where {T}
     if 0 <= atom.t <= 1
         add_constraint!(context, tmp in GeomMeanHypoCone(I, A, atom.t, false))
     else
-        add_constraint!(context, tmp in GeomMeanEpiCone(I, A, atom.t, false))
+        add_constraint!(
+            context,
+            Convex.GenericConstraint(
+                (tmp, I, A),
+                GeometricMeanEpiConeSquare(atom.t, size(A, 1)),
+            ),
+        )
     end
     # It's already a real mathematically, but Convex doesn't know it
     return conic_form!(context, real(LinearAlgebra.tr(atom.C * tmp)))

src/constraints/GeometricMeanEpiCone.jl

@@ -0,0 +1,102 @@
+# Copyright (c) 2014: Madeleine Udell and contributors
+# Copyright (c) 2021: Hamza Fawzi
+#
+# Use of this source code is governed by an MIT-style license that can be found
+# in the LICENSE.md file or at https://opensource.org/licenses/MIT.
+
+"""
+    GeometricMeanEpiConeSquare(t::Rational, side_dimension::Int)
+
+The constraint `(T, A, B) in GeometricMeanEpiConeSquare(t, side_dimension)`
+constrains T to
+
+  A #_t B ⪯ T
+
+where:
+
+  * A #_t B is the t-weighted geometric mean of A and B:
+    A^{1/2} (A^{-1/2} B A^{-1/2})^t A^{1/2}
+  * Parameter t must be in [-1, 0] or [1, 2].
+  * Constraints A ⪰ 0, B ⪰ 0 are added.
+
+## Reference
+
+Ported from CVXQUAD which is based on the paper: "Lieb's concavity theorem,
+matrix geometric means and semidefinite optimization" by Hamza Fawzi and James
+Saunderson (arXiv:1512.03401)
+"""
+struct GeometricMeanEpiConeSquare <: MOI.AbstractVectorSet
+    t::Rational
+    side_dimension::Int
+
+    function GeometricMeanEpiConeSquare(t::Rational, side_dimension::Int)
+        if !(-1 <= t <= 0 || 1 <= t <= 2)
+            throw(DomainError(t, "t must be in the range [-1, 0] or [1, 2]"))
+        end
+        return new(t, side_dimension)
+    end
+end
+
+MOI.dimension(set::GeometricMeanEpiConeSquare) = 3 * set.side_dimension^2
+
+function head(io::IO, ::GeometricMeanEpiConeSquare)
+    return print(io, "GeometricMeanEpiConeSquare")
+end
+
+function GenericConstraint(func::Tuple, set::GeometricMeanEpiConeSquare)
+    for f in func
+        n = LinearAlgebra.checksquare(f)
+        if n != set.side_dimension
+            throw(
+                DimensionMismatch(
+                    "Matrix of side dimension `$n` does not match set of side dimension `$(set.side_dimension)`",
+                ),
+            )
+        end
+    end
+    return GenericConstraint(vcat(vec.(func)...), set)
+end
+
+function _get_matrices(c::GenericConstraint{GeometricMeanEpiConeSquare})
+    n = c.set.side_dimension
+    d = n^2
+    T = reshape(c.child[1:d], n, n)
+    A = reshape(c.child[d.+(1:d)], n, n)
+    B = reshape(c.child[2d.+(1:d)], n, n)
+    return T, A, B
+end
+
+# For t ∈ [-1, 0] ∪ [1, 2] the t-weighted matrix geometric mean is matrix convex
+# (arxiv:1512.03401).
+# So if A and B are convex sets, then T ⪰ A #_t B will be a convex set.
+function vexity(constraint::GenericConstraint{GeometricMeanEpiConeSquare})
+    T, A, B = _get_matrices(constraint)
+    if vexity(A) in (ConcaveVexity(), NotDcp()) ||
+       vexity(B) in (ConcaveVexity(), NotDcp())
+        return NotDcp()
+    end
+    return ConvexVexity() + -vexity(T)
+end
+
+function _add_constraint!(
+    context::Context,
+    constraint::GenericConstraint{GeometricMeanEpiConeSquare},
+)
+    T, A, B = _get_matrices(constraint)
+    t = constraint.set.t
+    Z = if sign(constraint.child) == ComplexSign()
+        HermitianSemidefinite(size(A, 1))
+    else
+        Semidefinite(size(A, 1))
+    end
+    if t <= 0
+        add_constraint!(context, [T A; A Z] ⪰ 0)
+        add_constraint!(context, Z in GeomMeanHypoCone(A, B, -t, false))
+    else
+        # range of t checked in GeometricMeanEpiConeSquare constructor.
+        @assert t >= 1
+        add_constraint!(context, [T B; B Z] ⪰ 0)
+        add_constraint!(context, Z in GeomMeanHypoCone(A, B, 2 - t, false))
+    end
+    return
+end

src/expressions.jl

@@ -75,6 +75,8 @@ function vexity(x::AbstractExpr)
     return vex
 end
 
+vexity(::Value) = ConstVexity()
+
 evaluate(x) = output(x) # fallback
 
 Base.size(x::AbstractExpr) = x.size