Doc improvements for some atoms, plus an Expression.cast function. (cvxpy#2847)

* documentation and style for quantum_cond_entr

* documentation for von_neumann_entr and input checking thats consistent with other quantum information atoms

* documentation and evaluation tolerance for quantum_rel_entr

* remove unused function definitions;

* add a hermitian_wrap

* documentation, and create new Expression.cast function whose meaning is clearer than Expression.cast_to_const

* make ruff happy

Author: rileyjmurray

cvxpy/atoms/quantum_cond_entr.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Literal, Tuple
 
 import numpy as np
 
@@ -9,12 +9,54 @@
 from cvxpy.expressions.expression import Expression
 
 
-def quantum_cond_entr(rho: Expression , dim: list[int], sys: Optional[int]=0):
+def quantum_cond_entr(
+        rho: Expression , dims: Tuple[int, int], sys: Literal[0, 1] = 0, quad_approx=(3, 3)
+    ):
+    """
+    Returns (an approximation of) the quantum conditional entropy for a bipartite state,
+    conditioning on system :math:`\\texttt{sys}.`
+
+    Formally, if :math:`N` is the von Neumann entropy function and
+    :math:`\\operatorname{tr}_{\\texttt{sys}}` is the partial trace operator over subsystem
+    :math:`\\texttt{sys},` the returned expression represents
+
+    .. math::
+        N(\\rho) - N(\\operatorname{tr}_{\\texttt{sys}}(\\rho)).
+
+    Parameters
+    ----------
+    rho : Expression
+        A Hermitian matrix of order :math:`\\texttt{dims[0]}\\cdot\\texttt{dims[1]}.`
+
+    dims : tuple
+        The dimensions of the two subsystems that definte :math:`\\rho` as a bipartite state.
+
+    sys : int
+        The subsystem on which to condition in evaluating the conditional quantum entropy.
+
+    quad_approx : Tuple[int, int]
+        quad_approx[0] is the number of quadrature nodes and quad_approx[1] is the number of scaling
+        points in the quadrature scheme from https://arxiv.org/abs/1705.00812.
+
+    Notes
+    -----
+    This function does not assume :math:`\\operatorname{tr}(\rho)=1,` which would be required
+    for most uses of this function in the context of quantum information theory. See 
+    https://en.wikipedia.org/wiki/Conditional_quantum_entropy for more information. 
+    """
+    if len(dims) != 2:
+        err = 'This function is only defined for a tensor product of two subsystems,' + \
+        f'but {len(dims)} subsystems were implied from the value of dims.'
+        raise ValueError(err)
     if sys == 0:
-        composite_arg = kron(np.eye(dim[0]),
-                             partial_trace(rho, dim, sys))
-        return -quantum_rel_entr(rho, hermitian_wrap(composite_arg))
+        composite_arg = kron(
+            np.eye(dims[0]), partial_trace(rho, dims, sys)
+        )
     elif sys == 1:
-        composite_arg = kron(partial_trace(rho, dim, sys),
-                             np.eye(dim[1]))
-        return -quantum_rel_entr(rho, hermitian_wrap(composite_arg))
+        composite_arg = kron(
+            partial_trace(rho, dims, sys), np.eye(dims[1])
+        )
+    else:
+        raise ValueError(f'Argument sys must be either 0 or 1; got {sys}.')
+    composite_arg = hermitian_wrap(composite_arg)
+    return -quantum_rel_entr(rho, composite_arg, quad_approx)

cvxpy/atoms/quantum_rel_entr.py

@@ -19,22 +19,40 @@
 from scipy import linalg as LA
 from scipy.stats import entropy
 
+from cvxpy import settings
 from cvxpy.atoms.atom import Atom
 from cvxpy.constraints.constraint import Constraint
 
 
 class quantum_rel_entr(Atom):
     """
+    An approximation of the quantum relative entropy between systems with (possibly un-normalized) 
+    density matrices :math:`X` and :math`Y:`
+     
+    .. math::
+        \\operatorname{tr}\\left( X ( \\log X - \\log Y ) \\right).
+      
+    The approximation uses a quadrature scheme described in https://arxiv.org/abs/1705.00812.
+
     Parameters
     ----------
     X : Expression or numeric
         A PSD matrix
     Y : Expression or numeric
         A PSD matrix
+    quad_approx : Tuple[int, int]
+        quad_approx[0] is the number of quadrature nodes and quad_approx[1] is the number of scaling
+        points in the quadrature scheme from https://arxiv.org/abs/1705.00812.
+
+    Notes
+    -----
+    This function does not assume :math:`\\operatorname{tr}(X)=\\operatorname{tr}(Y)=1,` which
+    would be required for most uses of this function in the context of quantum information.
     """
 
+    EVAL_TOL = min(settings.ATOM_EVAL_TOL, 1e-6)
+
     def __init__(self, X, Y, quad_approx: Tuple[int, int] = (3, 3)) -> None:
-        # TODO: add a check that N is symmetric/Hermitian.
         self.quad_approx = quad_approx
         super(quantum_rel_entr, self).__init__(X, Y)
 
@@ -49,12 +67,11 @@ def numeric(self, values):
         w1, V = LA.eigh(X)
         w2, W = LA.eigh(Y)
         u = w1.T @ np.abs(V.conj().T @ W) ** 2
-        def func(x):
-            assert np.all(x >= 0)
-            x_sum = x.sum()
-            val = -entropy(x)
-            un_normalized = (x_sum * val + np.log(x_sum)*x_sum)
-            return un_normalized
+        if np.any(w1 < - self.EVAL_TOL) or np.any(w2 < -self.EVAL_TOL):
+            return np.inf
+        else:
+            w1[w1 < 0] = 0
+            w2[w2 < 0] = 0
         r1 = -entropy(w1)
         r2 = u @ np.log(w2)
         return (r1 - r2)

cvxpy/atoms/von_neumann_entr.py

@@ -25,32 +25,37 @@
 
 class von_neumann_entr(Atom):
     """
-    Computes the Von Neumann Entropy of the positive-definite matrix
-    :math:`X\\in\\mathbb{S}^n_{+}`
+    Represents the von Neumann Entropy of the positive-definite matrix :math:`X,`
 
         .. math::
-            -\\operatorname{tr}(X \\operatorname{logm}(X))
+            -\\operatorname{tr}(X \\log X).
 
-    where :math:`\\operatorname{tr}` is the trace and
-    :math:`\\operatorname{logm}` is the matrix logarithm
-
-    | May alternatively be expressed as:
+    Mathematically, this is equivalent to
 
         .. math::
             \\texttt{von_neumann_entr}(X) = -\\textstyle\\sum_{i=1}^n \\lambda_i \\log \\lambda_i
 
-    | where :math:`\\lambda_{i}` are the eigenvalues of :math:`X`
-    This atom does not enforce :math:`\\operatorname{tr}(X) = 1`
-    as is expected in applications from quantum mechanics.
+    where :math:`\\lambda_{i}` are the eigenvalues of :math:`X.`
 
     Parameters
     ----------
     X : Expression or numeric
         A PSD matrix
+
+    quad_approx : Tuple[int,...]
+        This is either an empty tuple (default) or a 2-tuple. If a 2-tuple, then this atom
+        is approximately canonicalized. The approximations replace ExpCone constraints with SOC
+        constraints based on a quadrature scheme from https://arxiv.org/abs/1705.00812;
+        quad_approx[0] is the number of quadrature nodes and quad_approx[1] is the number of
+        scaling points in the quadrature scheme.
+
+    Notes
+    -----
+    This function does not assume :math:`\\operatorname{tr}(X)=1,` which would be required
+    for most uses of this function in the context of quantum information theory. 
     """
 
-    def __init__(self, X, quad_approx: Tuple[int, int] = ()) -> None:
-        # TODO: add a check that N is symmetric/Hermitian.
+    def __init__(self, X, quad_approx: Tuple[int, ...] = ()) -> None:
         self.quad_approx = quad_approx
         super(von_neumann_entr, self).__init__(X)
 
@@ -63,10 +68,14 @@ def numeric(self, values):
         return val
 
     def validate_arguments(self) -> None:
-        """Verify that the argument is a square matrix."""
-        if not self.args[0].ndim == 2 or self.args[0].shape[0] != self.args[0].shape[1]:
+        """Verify that the argument is Hermitian."""
+        if not self.args[0].is_hermitian():
             raise ValueError(
-                f"The argument {self.args[0].name()} to von_neumann_entr must be a 2-d square array"
+                f"""
+                The argument {self.args[0].name()} to von_neumann_entr must be a Hermitian matrix.
+                If you know for a fact that the input is Hermitian, wrap it with the hermitian_wrap
+                atom before calling von_neumann_entr.
+                """
             )
 
     def sign_from_args(self) -> Tuple[bool, bool]:
@@ -118,7 +127,7 @@ def _grad(self, values):
         # derivative = 2*(L + L * logm(np.dot(L.T, L)))
         # TODO: have to wrap derivative around scipy CSC sparse matrices
         #  compare to log_det atom.
-        raise ValueError()
+        raise NotImplementedError()
 
     def _domain(self) -> List[Constraint]:
         """Returns constraints describing the domain of the node.

cvxpy/constraints/exponential.py

@@ -58,9 +58,9 @@ class ExpCone(Cone):
 
     def __init__(self, x: Expression, y: Expression, z: Expression, constr_id=None) -> None:
         Expression = cvxtypes.expression()
-        self.x = Expression.cast_to_const(x)
-        self.y = Expression.cast_to_const(y)
-        self.z = Expression.cast_to_const(z)
+        self.x = Expression.cast(x)
+        self.y = Expression.cast(y)
+        self.z = Expression.cast(z)
         args = [self.x, self.y, self.z]
         for val in args:
             if not (val.is_affine() and val.is_real()):
@@ -163,19 +163,15 @@ def _dual_cone(self, *args):
 
 
 class RelEntrConeQuad(Cone):
-    """An approximate construction of the scalar relative entropy cone
-
-    Definition:
+    """An approximation of the scalar relative entropy cone,
 
     .. math::
 
         K_{re}=\\text{cl}\\{(x,y,z)\\in\\mathbb{R}_{++}\\times
-                \\mathbb{R}_{++}\\times\\mathbb{R}_{++}\\:x\\log(x/y)\\leq z\\}
-
-    Since the above definition is very similar to the ExpCone, we provide a conversion method.
+                \\mathbb{R}_{++}\\times\\mathbb{R}_{++}\\:x\\log(x/y)\\leq z\\},
 
-    More details on the approximation can be found in Theorem-3 on page-10 in the paper:
-    Semidefinite Approximations of the Matrix Logarithm.
+    in terms of second order cones. The approximation uses a numerical quadrature scheme
+    described in https://arxiv.org/abs/1705.00812.
 
     Parameters
     ----------
@@ -185,17 +181,18 @@ class RelEntrConeQuad(Cone):
         y in the (approximate) scalar relative entropy cone
     z : Expression
         z in the (approximate) scalar relative entropy cone
-    m: Parameter directly related to the number of generated nodes for the quadrature
-    approximation used in the algorithm
-    k: Another parameter controlling the approximation
+    m : int
+        Number of quadrature points in the approximation.
+    k: int
+        Number of scaling points in the approximation.
     """
 
     def __init__(self, x: Expression, y: Expression, z: Expression,
                  m: int, k: int, constr_id=None) -> None:
         Expression = cvxtypes.expression()
-        self.x = Expression.cast_to_const(x)
-        self.y = Expression.cast_to_const(y)
-        self.z = Expression.cast_to_const(z)
+        self.x = Expression.cast(x)
+        self.y = Expression.cast(y)
+        self.z = Expression.cast(z)
         args = [self.x, self.y, self.z]
         for val in args:
             if not (val.is_affine() and val.is_real()):
@@ -281,17 +278,15 @@ def save_dual_value(self, value) -> None:
 
 
 class OpRelEntrConeQuad(Cone):
-    """An approximate construction of the operator relative entropy cone
-
-    Definition:
+    """An approximate construction of the operator relative entropy cone,
 
     .. math::
 
-        K_{re}^n=\\text{cl}\\{(X,Y,T)\\in\\mathbb{H}^n_{++}\\times
-                \\mathbb{H}^n_{++}\\times\\mathbb{H}^n_{++}\\:D_{\\text{op}}\\succeq T\\}
+        K_{re}^n = \\text{cl}\\{(X,Y,T)\\in\\mathbb{H}^n_{++}\\times
+                \\mathbb{H}^n_{++}\\times\\mathbb{H}^n_{++}\\:D_{\\text{op}}(X,Y) \\succeq T\\}.
 
-    More details on the approximation can be found in Theorem-3 on page-10 in the paper:
-    Semidefinite Approximations of the Matrix Logarithm.
+    Details on the approximation can be found in Theorem-3 on page-10 of
+    https://arxiv.org/abs/1705.00812.
 
     Parameters
     ----------
@@ -317,9 +312,9 @@ class OpRelEntrConeQuad(Cone):
     def __init__(self, X: Expression, Y: Expression, Z: Expression,
                  m: int, k: int, constr_id=None) -> None:
         Expression = cvxtypes.expression()
-        self.X = Expression.cast_to_const(X)
-        self.Y = Expression.cast_to_const(Y)
-        self.Z = Expression.cast_to_const(Z)
+        self.X = Expression.cast(X)
+        self.Y = Expression.cast(Y)
+        self.Z = Expression.cast(Z)
         if (not X.is_hermitian()) or (not Y.is_hermitian()) or (not Z.is_hermitian()):
             msg = ("One of the input matrices has not explicitly been declared as symmetric or"
                    "Hermitian. If the inputs are Variable objects, try declaring them with the"

cvxpy/expressions/expression.py

@@ -564,6 +564,15 @@ def __rpow__(self, base: float) -> "Expression":
                                   " identity that a**x = cp.exp(cp.multiply(np"
                                   ".log(a), x)).")
 
+    @staticmethod
+    def cast(expr_like) -> "Expression":
+        """
+        If expr_like is an Expression, return it. Otherwise, cast expr_like to a Constant.
+
+        This is a wrapper around the misleadingly-named `Expression.cast_to_const` function.
+        """
+        return Expression.cast_to_const(expr_like)
+
     # Arithmetic operators.
     @staticmethod
     def cast_to_const(expr: "Expression"):

cvxpy/tests/test_quantum_rel_entr.py

@@ -11,7 +11,7 @@
 from cvxpy.tests import solver_test_helpers as STH
 
 
-def applychan(chan: np.array, rho: cp.Variable, rep: str, dim: tuple[int, int]):
+def applychan(chan: np.ndarray, rho: cp.Variable, rep: str, dim: tuple[int, int]):
     dimA, dimB, dimE = None, None, None
     if rep == 'choi2':
         dimA, dimB = dim
@@ -25,14 +25,6 @@ def applychan(chan: np.array, rho: cp.Variable, rep: str, dim: tuple[int, int]):
         rho_out = partial_trace(chan @ rho @ chan.conj().T, [dimB, dimE], 1)
         return rho_out
 
-def randH(n: int):
-    A = np.random.randn(n, n) + 1j * np.random.randn(n, n)
-    return (A + A.conj().T)/2
-
-def randRho(n: int):
-    p = 10 * randH(n)
-    p = (p @ p.conj().T)/np.trace(p @ p.conj().T)
-    return p
 
 class TestQuantumRelEntr:
     """

cvxpy/tests/test_von_neumann_entr.py

@@ -262,8 +262,9 @@ def make_test_4():
         p2_expect = 0.5
         var_pairs = [(p1, p1_expect), (p2, p2_expect)]
 
-        obj = cp.Maximize((cp.von_neumann_entr(p1 * rho1 + p2 * rho2) - p1 * H1 - \
-                           p2 * H2)/np.log(2))
+        vne_arg = cp.hermitian_wrap(p1 * rho1 + p2 * rho2)
+        obj_arg = (cp.von_neumann_entr(vne_arg) - p1 * H1 - p2 * H2)/np.log(2)
+        obj = cp.Maximize(obj_arg)
         obj_expect = 0.60088
         obj_pair = (obj, obj_expect)