Add utilities for constructing CVXPY models involving (convex) diamon…

…d-norm expressions or (concave) root-fidelity expressions. Add a "normalize" keyword argument to the various functions for computing(/uncomputing) Jamiolkowski isomorphisms of operation matrices. Added option to return all optimization model variables used in computing diamonddist.

pygsti/tools/jamiolkowski.py

@@ -64,9 +64,10 @@
 #  J(Phi) = sum_(0<i,j<n) Phi(Eij) otimes Eij                                                                                                                   # noqa
 #  where Eij is the matrix unit with a single element in the (i,j)-th position, i.e. Eij == |i><j|                                                              # noqa
 
-def jamiolkowski_iso(operation_mx, op_mx_basis='pp', choi_mx_basis='pp'):
+def jamiolkowski_iso(operation_mx, op_mx_basis='pp', choi_mx_basis='pp', normalized=True):
     """
-    Given a operation matrix, return the corresponding Choi matrix that is normalized to have trace == 1.
+    Return a Choi matrix (in the choi_mx_basis) for operation_mx, when operation_mx
+    is interpreted in the op_mx_basis. 
 
     Parameters
     ----------
@@ -83,12 +84,23 @@ def jamiolkowski_iso(operation_mx, op_mx_basis='pp', choi_mx_basis='pp'):
         values are Matrix-unit (std), Gell-Mann (gm), Pauli-product (pp),
         and Qutrit (qt) (or a custom basis object).
 
+    normalized : bool
+        If normalized=True, then this function maps trace-preserving 
+        operation matrices to trace-1 Choi matrices.
+
     Returns
     -------
     numpy array
-        the Choi matrix, normalized to have trace == 1, in the desired basis.
+        the Choi matrix, in the desired basis.
     """
-    operation_mx = _np.asarray(operation_mx)
+    try:
+        import cvxpy as cp
+        is_cvxpy_expression = isinstance(operation_mx, cp.Expression)
+    except ImportError:
+        is_cvxpy_expression = False
+
+    if not is_cvxpy_expression:
+        operation_mx = _np.asarray(operation_mx)
     op_mx_basis = _bt.create_basis_for_matrix(operation_mx, op_mx_basis)
     opMxInStdBasis = _bt.change_basis(operation_mx, op_mx_basis, op_mx_basis.create_equivalent('std'))
 
@@ -113,25 +125,33 @@ def jamiolkowski_iso(operation_mx, op_mx_basis='pp', choi_mx_basis='pp'):
     M = len(BVec)  # can be < N if basis has multiple block dims
     assert(M == N), 'Expected {}, got {}'.format(M, N)
 
-    choiMx = _np.empty((N, N), 'complex')
+    opMxInStdBasis_vec = opMxInStdBasis.reshape((N*N, 1))
+    # ^ An explicit column vector.
+    choiMx_cols = []
     for i in range(M):
+        rows = []
         for j in range(M):
             BiBj = _np.kron(BVec[i], _np.conjugate(BVec[j]))
-            num = _np.vdot(BiBj, opMxInStdBasis)
-            den = _np.linalg.norm(BiBj) ** 2
-            choiMx[i, j] = num / den 
-
+            BiBj /= _np.linalg.norm(BiBj) ** 2
+            rows.append(BiBj.conj().ravel())
+        choiMx_cols.append( _np.array(rows) @ opMxInStdBasis_vec )
+    if is_cvxpy_expression:
+        choiMx = cp.hstack(choiMx_cols)
+    else:
+        choiMx = _np.hstack(choiMx_cols)
     # This construction results in a Jmx with trace == dim(H) = sqrt(operation_mx.shape[0])
     #  (dimension of density matrix) but we'd like a Jmx with trace == 1, so normalize:
-    choiMx_normalized = choiMx / dmDim
-    return choiMx_normalized
+    if normalized:
+        choiMx /= dmDim
+    return choiMx
 
 # GStd = sum_ij Jij (BSi x BSj^*)
 
 
-def jamiolkowski_iso_inv(choi_mx, choi_mx_basis='pp', op_mx_basis='pp'):
+def jamiolkowski_iso_inv(choi_mx, choi_mx_basis='pp', op_mx_basis='pp', normalized=True):
     """
-    Given a choi matrix, return the corresponding operation matrix.
+    Given a choi matrix (interpreted in choi_mx_basis), return the corresponding
+    operation matrix (in op_mx_basis).
 
     This function performs the inverse of :func:`jamiolkowski_iso`.
 
@@ -150,6 +170,10 @@ def jamiolkowski_iso_inv(choi_mx, choi_mx_basis='pp', op_mx_basis='pp'):
         values are Matrix-unit (std), Gell-Mann (gm), Pauli-product (pp),
         and Qutrit (qt) (or a custom basis object).
 
+    normalized : bool
+        If normalized=True, then we assume choi_mx was computed with the 
+        convention that trace-preserving maps have trace-1 choi matrices.
+
     Returns
     -------
     numpy array
@@ -167,7 +191,10 @@ def jamiolkowski_iso_inv(choi_mx, choi_mx_basis='pp', op_mx_basis='pp'):
     assert(len(BVec) == N)  # make sure the number of basis matrices matches the dim of the choi matrix given
 
     # Invert normalization
-    choiMx_unnorm = choi_mx * dmDim
+    if normalized:
+        choiMx_unnorm = choi_mx * dmDim
+    else:
+        choiMx_unnorm = choi_mx
 
     opMxInStdBasis = _np.zeros((N, N), 'complex')  # in matrix unit basis of entire density matrix
     for i in range(N):
@@ -187,9 +214,10 @@ def jamiolkowski_iso_inv(choi_mx, choi_mx_basis='pp', op_mx_basis='pp'):
     return _bt.change_basis(opMxInStdBasis, op_mx_basis.create_equivalent('std'), op_mx_basis)
 
 
-def fast_jamiolkowski_iso_std(operation_mx, op_mx_basis):
+def fast_jamiolkowski_iso_std(operation_mx, op_mx_basis, normalized=True):
     """
-    The corresponding Choi matrix in the standard basis that is normalized to have trace == 1.
+    Returns the standard-basis representation of the Choi matrix for operation_mx,
+    where operation_mx is interpreted in op_mx_basis.
 
     This routine *only* computes the case of the Choi matrix being in the
     standard (matrix unit) basis, but does so more quickly than
@@ -206,6 +234,10 @@ def fast_jamiolkowski_iso_std(operation_mx, op_mx_basis):
         values are Matrix-unit (std), Gell-Mann (gm), Pauli-product (pp),
         and Qutrit (qt) (or a custom basis object).
 
+    normalized : bool
+        If normalized=True, then this function maps trace-preserving 
+        operation matrices to trace-1 Choi matrices. 
+
     Returns
     -------
     numpy array
@@ -232,13 +264,15 @@ def fast_jamiolkowski_iso_std(operation_mx, op_mx_basis):
 
     # This construction results in a Jmx with trace == dim(H) = sqrt(gateMxInPauliBasis.shape[0])
     #  but we'd like a Jmx with trace == 1, so normalize:
-    Jmx_norm = Jmx / N
-    return Jmx_norm
+    if normalized:
+        Jmx /= N
+    return Jmx
 
 
-def fast_jamiolkowski_iso_std_inv(choi_mx, op_mx_basis):
+def fast_jamiolkowski_iso_std_inv(choi_mx, op_mx_basis, normalized=True):
     """
-    Given a choi matrix in the standard basis, return the corresponding operation matrix.
+    Given a choi matrix in the standard basis, return the corresponding
+    operation matrix (in op_mx_basis).
 
     This function performs the inverse of :func:`fast_jamiolkowski_iso_std`.
 
@@ -253,6 +287,10 @@ def fast_jamiolkowski_iso_std_inv(choi_mx, op_mx_basis):
         values are Matrix-unit (std), Gell-Mann (gm), Pauli-product (pp),
         and Qutrit (qt) (or a custom basis object).
 
+    normalized : bool
+        If normalized=True, then we assume choi_mx was computed with the 
+        convention that trace-preserving maps have trace-1 choi matrices.
+
     Returns
     -------
     numpy array
@@ -262,7 +300,9 @@ def fast_jamiolkowski_iso_std_inv(choi_mx, op_mx_basis):
     #Shuffle indices to go from process matrix to Jamiolkowski matrix (they vectorize differently)
     N2 = choi_mx.shape[0]; N = int(_np.sqrt(N2))
     assert(N * N == N2)  # make sure N2 is a perfect square
-    opMxInStdBasis = choi_mx.reshape((N, N, N, N)) * N
+    opMxInStdBasis = choi_mx.reshape((N, N, N, N))
+    if normalized:
+        opMxInStdBasis *= N
     opMxInStdBasis = _np.swapaxes(opMxInStdBasis, 1, 2).ravel()
     opMxInStdBasis = opMxInStdBasis.reshape((N2, N2))
     op_mx_basis = _bt.create_basis_for_matrix(opMxInStdBasis, op_mx_basis)

pygsti/tools/optools.py

@@ -10,6 +10,11 @@
 # http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
 #***************************************************************************************************
 
+from __future__ import annotations
+from typing import TYPE_CHECKING, List, Tuple
+if TYPE_CHECKING:
+    import cvxpy as cp
+
 import collections as _collections
 import warnings as _warnings
 
@@ -23,6 +28,7 @@
 from pygsti.tools import jamiolkowski as _jam
 from pygsti.tools import lindbladtools as _lt
 from pygsti.tools import matrixtools as _mt
+from pygsti.tools import sdptools as _sdps
 from pygsti.baseobjs import basis as _pgb
 from pygsti.baseobjs.basis import Basis as _Basis, ExplicitBasis as _ExplicitBasis, DirectSumBasis as _DirectSumBasis
 from pygsti.baseobjs.label import Label as _Label
@@ -270,7 +276,7 @@ def tracedist(a, b):
     return 0.5 * tracenorm(a - b)
 
 
-def diamonddist(a, b, mx_basis='pp', return_x=False):
+def diamonddist(a, b, mx_basis='pp', return_x=False, return_all_vars=False):
     """
     Returns the approximate diamond norm describing the difference between gate matrices.
 
@@ -303,98 +309,32 @@ def diamonddist(a, b, mx_basis='pp', return_x=False):
         Only returned if `return_x = True`.  Encodes the state rho, such that
         `dm = trace( |(J(a)-J(b)).T * W| )`.
     """
+    assert _sdps.CVXPY_ENABLED
+    assert a.shape == b.shape
     mx_basis = _bt.create_basis_for_matrix(a, mx_basis)
 
-    # currently cvxpy is only needed for this function, so don't import until here
-    import cvxpy as _cvxpy
-    old_cvxpy = bool(tuple(map(int, _cvxpy.__version__.split('.'))) < (1, 0))
-    if old_cvxpy:
-        raise RuntimeError('CVXPY 0.4 is no longer supported. Please upgrade to CVXPY 1.0 or higher.')
-
+    c = b - a
     # _jam code below assumes *un-normalized* Jamiol-isomorphism.
-    # It will convert a & b to a "single-block" basis representation
-    # when mx_basis has multiple blocks. So after we call it, we need
-    # to multiply by mx dimension (`smallDim`).
-    JAstd = _jam.fast_jamiolkowski_iso_std(a, mx_basis)
-    JBstd = _jam.fast_jamiolkowski_iso_std(b, mx_basis)
-    dim = JAstd.shape[0]
-    smallDim = int(_np.sqrt(dim))
-    JAstd *= smallDim
-    JBstd *= smallDim
-    assert(dim == JAstd.shape[1] == JBstd.shape[0] == JBstd.shape[1])
-
-    J = JBstd - JAstd
-    prob, vars = _diamond_norm_model(dim, smallDim, J)
+    # It will convert c a "single-block" basis representation
+    # when mx_basis has multiple blocks.
+    J = _jam.fast_jamiolkowski_iso_std(c, mx_basis, normalized=False)
+    prob, vars = _sdps.diamond_norm_model_jamiolkowski(J)
 
     try:
         prob.solve(solver='CVXOPT')
-    except _cvxpy.error.SolverError as e:
+        objective_val = prob.value
+        varvals = [v.value for v in vars]
+    except Exception as e:
         _warnings.warn("CVXPY failed: %s - diamonddist returning -2!" % str(e))
-        return (-2, _np.zeros((dim, dim))) if return_x else -2
-    except:
-        _warnings.warn("CVXOPT failed (unknown err) - diamonddist returning -2!")
-        return (-2, _np.zeros((dim, dim))) if return_x else -2
-
-    if return_x:
-        return prob.value, vars[0].value
-    else:
-        return prob.value
+        objective_val = -2
+        varvals = [_np.zeros_like(J), None, None]
 
-
-def _diamond_norm_model(dim, smallDim, J):
-    # return a model for computing the diamond norm.
-    #
-    # Uses the primal SDP from arXiv:1207.5726v2, Sec 3.2
-    #
-    # Maximize 1/2 ( < J(phi), X > + < J(phi).dag, X.dag > )
-    # Subject to  [[ I otimes rho0,       X        ],
-    #              [      X.dag   ,   I otimes rho1]] >> 0
-    #              rho0, rho1 are density matrices
-    #              X is linear operator
-
-    import cvxpy as _cp
-
-    rho0 = _cp.Variable((smallDim, smallDim), name='rho0', hermitian=True)
-    rho1 = _cp.Variable((smallDim, smallDim), name='rho1', hermitian=True)
-    X = _cp.Variable((dim, dim), name='X', complex=True)
-    Y = _cp.real(X)
-    Z = _cp.imag(X)
-
-    K = J.real
-    L = J.imag
-    if hasattr(_cp, 'scalar_product'):
-        objective_expr = _cp.scalar_product(K, Y) + _cp.scalar_product(L, Z)
+    if return_all_vars:
+        return objective_val, varvals
+    elif return_x:
+        return objective_val, varvals[0]
     else:
-        Kf = K.flatten(order='F')
-        Yf = Y.flatten(order='F')
-        Lf = L.flatten(order='F')
-        Zf = Z.flatten(order='F')
-        objective_expr = Kf @ Yf + Lf @ Zf
-
-    objective = _cp.Maximize(objective_expr)
-
-    ident = _np.identity(smallDim, 'd')
-    kr_tau0 = _cp.kron(ident, _cp.imag(rho0))
-    kr_tau1 = _cp.kron(ident, _cp.imag(rho1))
-    kr_sig0 = _cp.kron(ident, _cp.real(rho0))
-    kr_sig1 = _cp.kron(ident, _cp.real(rho1))
-
-    block_11 = _cp.bmat([[kr_sig0 ,    Y   ],
-                         [   Y.T  , kr_sig1]])
-    block_21 = _cp.bmat([[kr_tau0 ,    Z   ],
-                         [   -Z.T , kr_tau1]])
-    block_12 = block_21.T
-    mat_joint = _cp.bmat([[block_11, block_12],
-                          [block_21, block_11]])
-    constraints = [
-        mat_joint >> 0,
-        rho0 >> 0,
-        rho1 >> 0,
-        _cp.trace(rho0) == 1.,
-        _cp.trace(rho1) == 1.
-    ]
-    prob = _cp.Problem(objective, constraints)
-    return prob, [X, rho0, rho1]
+        return objective_val
 
 
 def jtracedist(a, b, mx_basis='pp'):  # Jamiolkowski trace distance:  Tr(|J(a)-J(b)|)