Refactoring gaugeopt.py: maintainability and fidelity as an objective

pygsti/models/explicitmodel.py

@@ -10,6 +10,8 @@
 # http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
 #***************************************************************************************************
 
+from __future__ import annotations
+
 import collections as _collections
 import itertools as _itertools
 import uuid as _uuid
@@ -1635,6 +1637,10 @@ def add_availability(opkey, op):
                                        availability,
                                        instrument_names=list(self.instruments.keys()), nonstd_instruments=self.instruments)
 
+    def copy(self) -> ExplicitOpModel:
+        c = super().copy()  # <-- that is indeed an ExplicitOpModel
+        return c # type: ignore
+
     def create_modelmember_graph(self):
         return _MMGraph({
             'preps': self.preps,

pygsti/models/model.py

@@ -10,6 +10,8 @@
 # http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
 #***************************************************************************************************
 
+from __future__ import annotations
+
 import bisect as _bisect
 import copy as _copy
 import itertools as _itertools
@@ -326,7 +328,7 @@ def _post_copy(self, copy_into, memo):
         """
         pass
 
-    def copy(self):
+    def copy(self) -> Model:
         """
         Copy this model.
 
@@ -2309,7 +2311,7 @@ def _post_copy(self, copy_into, memo):
         copy_into._reinit_opcaches()
         super(OpModel, self)._post_copy(copy_into, memo)
 
-    def copy(self):
+    def copy(self) -> OpModel:
         """
         Copy this model.
 
@@ -2320,7 +2322,7 @@ def copy(self):
         """
         self._clean_paramvec()  # ensure _paramvec is rebuilt if needed
         if OpModel._pcheck: self._check_paramvec()
-        ret = Model.copy(self)
+        ret = Model.copy(self) # <-- don't be fooled. That's an OpModel!
         if self._param_bounds is not None and self.parameter_labels is not None:
             ret._clean_paramvec()  # will *always* rebuild paramvec; do now so we can preserve param bounds
             assert _np.all(self.parameter_labels == ret.parameter_labels)  # ensure ordering is the same

pygsti/report/reportables.py

@@ -17,6 +17,7 @@
 """
 import importlib
 import warnings as _warnings
+from typing import Union
 
 import numpy as _np
 import scipy.linalg as _spl
@@ -40,6 +41,8 @@
 
 FINITE_DIFF_EPS = 1e-7
 
+BasisLike = Union[_Basis, str]
+
 
 def _null_fn(*arg):
     return None
@@ -1577,32 +1580,67 @@ def eigenvalue_nonunitary_avg_gate_infidelity(a, b, mx_basis):
 # init args == (model1, model2, op_label)
 
 
-def eigenvalue_entanglement_infidelity(a, b, mx_basis):
+
+def eigenvalue_entanglement_infidelity(a: _np.ndarray, b: _np.ndarray, mx_basis: BasisLike, is_tp=None, is_unitary=None, tol=1e-8) -> _np.floating:
     """
-    Eigenvalue entanglement infidelity between a and b
+    Eigenvalue entanglement infidelity is the infidelity between certain diagonal matrices that
+    contain [see Note 1] the eigenvalues of J(a) and J(b), where J(.) is the Jamiolkowski
+    isomorphism map. This is equivalent to computing infidelity of (J(a), J(b)) while pretending
+    that they share an eigenbasis.
 
     Parameters
     ----------
-    a : numpy.ndarray
-        The first process (transfer) matrix.
+    is_tp : bool, optional (default None)
+        Flag indicating that a and b are TP in their provided basis. If None (the default),
+        an explicit check is performed up to numerical tolerance `tol`. If True/False, then
+        the check is skipped and this is used as the result of the check as though it were 
+        performed.
 
-    b : numpy.ndarray
-        The second process (transfer) matrix.
+    is_unitary : bool, optional (default None)
+        Flag indicating that b is unitary. If None (the default) an explicit check is performed
+        up to tolerance `tol`. If True/False, then the check is skipped and this is used as the
+        result of the check as though it were performed.
 
-    mx_basis : Basis or {'pp', 'gm', 'std'}
-        the basis that `a` and `b` are in.
+    Notes
+    -----
+    [1] The eigenvalues are ordered using a min-weight matching algorithm.
 
-    Returns
-    -------
-    float
+    [2] If a and b are trace-preserving (TP) and b is unitary, then we compute eigenvalue
+        entanglement fidelity by leveraging three facts:
+
+        1. If (x,y) share an eigenbasis and have (consistently ordered) eigenvalues (v(x),v(y)),
+           then Tr(x y^{\\dagger}) is the inner product of v(x) and v(y).
+
+        2. J(a) and J(b) share an eigenbasis of if `a` and `b` share an eigenbasis.
+
+        3. If `a` and `b` are TP and `b` is unitary, then their entanglement fidelity can be
+           expressed as |Tr( a.b^{\\dagger} )| / d^2.
+
+       Then together, we see that if (a,b) are TP and `b` is unitary, then their eigenvalue 
+       entanglement fidelity is expressible as the inner product of v(a) and v(b), where 
+       v(a) and v(b) vectors holding are suitably-ordered eigenvalues of a and b.
+
     """
     d2 = a.shape[0]
-    evA = _np.linalg.eigvals(a)
-    evB = _np.linalg.eigvals(b)
-    _, pairs = _tools.minweight_match(evA, evB, lambda x, y: abs(x - y),
-                                      return_pairs=True)  # just to get pairing
-    mlPl = abs(_np.sum([_np.conjugate(evB[j]) * evA[i] for i, j in pairs]))
-    return 1.0 - mlPl / float(d2)
+
+    if is_unitary is None:
+        is_unitary = _np.allclose(_np.eye(d2), b @ b.T.conj(), atol=tol, rtol=tol)
+    if is_tp is None:
+        is_tp = _tools.is_trace_preserving(a, mx_basis, tol) and _tools.is_trace_preserving(b, mx_basis, tol) 
+
+    if is_unitary and is_tp:
+        evA = _np.linalg.eigvals(a)
+        evB = _np.linalg.eigvals(b)
+        _, pairs = _tools.minweight_match(evA, evB, lambda x, y: abs(x - y),
+                                        return_pairs=True)  # just to get pairing
+        fid = abs(_np.sum([_np.conjugate(evB[j]) * evA[i] for i, j in pairs])) / d2
+    else:
+        Ja = _tools.fast_jamiolkowski_iso_std(a, mx_basis)
+        Jb = _tools.fast_jamiolkowski_iso_std(b, mx_basis)
+        from pygsti.tools.optools import eigenvalue_fidelity
+        fid = eigenvalue_fidelity(Ja, Jb, gauge_invariant=True)
+    return 1.0 - fid
+
 
 
 Eigenvalue_entanglement_infidelity = _modf.opsfn_factory(eigenvalue_entanglement_infidelity)

pygsti/tools/matrixtools.py

@@ -90,6 +90,16 @@ def is_hermitian(mx, tol=1e-9):
         return _np.all(_np.abs(mx - mx.T.conj()) <= tol)
 
 
+def assert_hermitian(mat : _np.ndarray, tol: _np.floating) -> None:
+    hermiticity_error = _np.abs(mat - mat.T.conj())
+    if _np.any(hermiticity_error > tol):
+        message = f"""
+            Input matrix 'mat' is not Hermitian, up to tolerance {tol}.
+            The absolute values of entries in (mat - mat^H) are \n{hermiticity_error}. 
+        """
+        raise ValueError(message)
+
+
 def is_pos_def(mx, tol=1e-9, attempt_cholesky=False):
     """
     Test whether mx is a positive-definite matrix.

pygsti/tools/optools.py

@@ -184,6 +184,67 @@ def psd_square_root(mat):
     return f
 
 
+def eigenvalue_fidelity(x: _np.ndarray, y: _np.ndarray, gauge_invariant:bool=True) -> _np.floating:
+    """
+    For positive semidefinite matrices (x, y), define M(x,y) = √(√x · y · √x).
+
+    The fidelity of (x, y) is defined as F(x,y) = Tr(M(x,y))^2.
+
+    Let v(x) and v(y) denote vectors containing the eigenvalues of x and y,
+    sorted in decreasing order.
+
+    If gauge_invariant is True, then this function returns ⟨√v(x), √v(y)⟩^2,
+    which is always an upper bound for F(x,y).
+
+    If gauge_invariant is False, then this function uses a heuristic to "match"
+    eigenvalues of x and y based on similarity of their corresponding eigenvectors.
+
+    Proof
+    -----
+    Let A = √x·√y, so that M(x,y) = √(A A^†). The fact that M(x,y) is a square-root
+    of a Gram matrix of A implies (via the polar decomposition) the existence of a
+    unitary matrix U where A = M(x,y)U. This lets us write fidelity as
+
+        F(x,y) = Tr( √x · √y U^† )^2.                            (1)
+
+    Next, let s(z) denote the vector of singular values of a matrix z, in descending
+    order. Assuming the claim from Problem III.6.12 of Bhatia's Matrix Analysis, we
+    have the inequality
+
+        Tr( √x · √y U^† ) ≤ ⟨ s(√x), s(√y U^†) ⟩.                (2)
+
+    Using the fact that the singular values of √x are the same as its eigenvalues,
+    and the singular values of √y U^† are the eigenvalues of √y, we have
+
+        ⟨ s(√x), s(√y U^†) ⟩ = ⟨ √v(x), √v(y) ⟩.                 (3)
+
+    Combining (1), (2), and (3) proves our claim.
+    """
+    tol = _np.finfo(x.dtype).eps ** 0.75
+    _mt.assert_hermitian(x, tol)
+    _mt.assert_hermitian(y, tol)
+    if gauge_invariant:
+        valsX = _np.sort(_spl.eigvalsh(x))
+        valsY = _np.sort(_spl.eigvalsh(y))
+        pairs = [(i,i) for i in range(valsX.size)]
+    else:
+        from pygsti.tools import minweight_match
+        valsX, vecsX = _spl.eigh(x)
+        valsY, vecsY = _spl.eigh(y)
+        dissimilarity = lambda vec_x, vec_y : abs(1 - abs(vec_x @ vec_y))
+        _, pairs = minweight_match(vecsX.T.conj(), vecsY.T.conj(), dissimilarity, return_pairs=True)
+    ind_x, ind_y = zip(*pairs)
+    arg_x = _np.maximum(valsX[list(ind_x)], 0)
+    arg_y = _np.maximum(valsY[list(ind_y)], 0)
+    f = (arg_x**0.5 @ arg_y**0.5)**2
+    return f # type: ignore
+
+
+def eigenvalue_infidelity(a, b, gauge_invariant=True) -> _np.floating:
+    return 1 - eigenvalue_fidelity(a, b, gauge_invariant)
+
+
+
 def frobeniusdist(a, b) -> _np.floating[Any]:
     """
     Returns the frobenius distance between arrays: ||a - b||_Fro.
@@ -377,6 +438,26 @@ def jtracedist(a, b, mx_basis='pp'):  # Jamiolkowski trace distance:  Tr(|J(a)-J
     return tracedist(JA, JB)
 
 
+def is_trace_preserving(a: _np.ndarray, mx_basis: BasisLike='pp', tol=1e-8) -> bool:
+    dim   = int(_np.round( a.size**0.5 ))
+    udim  = int(_np.round( dim**0.5    ))
+    basis = _Basis.cast(mx_basis, dim=dim)
+    if basis.first_element_is_identity:
+        checkone = _np.isclose(a[0,0], 1.0, rtol=tol, atol=tol)
+        checktwo = _np.allclose(a[1:], 0.0, atol=tol)
+        return checkone and checktwo
+    # else, check that the adjoint of a is unital.
+    I_mat = _np.eye(udim)
+    I_vec = _bt.stdmx_to_vec(I_mat, basis)
+    if _np.isrealobj(a):
+        expect_I_vec = a.T @ I_vec
+    else:
+        expect_I_vec = a.T.conj() @ I_vec
+    atol = tol * udim
+    check = float(_np.linalg.norm(I_vec - expect_I_vec)) <= atol
+    return check
+
+
 def entanglement_fidelity(a, b, mx_basis: BasisLike='pp', is_tp=None, is_unitary=None):
     """
     Returns the "entanglement" process fidelity between gate  matrices.
@@ -439,10 +520,7 @@ def entanglement_fidelity(a, b, mx_basis: BasisLike='pp', is_tp=None, is_unitary
 
     #if the tp flag isn't set we'll calculate whether it is true here
     if is_tp is None:
-        def is_tp_fn(x):
-            return _np.isclose(x[0, 0], 1.0) and _np.allclose(x[0, 1:d2], 0)
-
-        is_tp= (is_tp_fn(a) and is_tp_fn(b))
+        is_tp = is_trace_preserving(a, mx_basis) and is_trace_preserving(b, mx_basis)
 
     #if the unitary flag isn't set we'll calculate whether it is true here 
     if is_unitary is None:

test/unit/algorithms/test_gaugeopt_correctness.py

@@ -270,7 +270,7 @@ def test_lbfgs_frodist(self):
 
     def test_lbfgs_tracedist(self):
         times = []
-        for seed in range(1, 3):
+        for seed in [1]:
             dt = self._main_tester('L-BFGS-B', seed, test_tol=1e-6, alg_tol=1e-8, gop_objective='tracedist')
             times.append(dt)
         print(f'L-BFGS GaugeOpt over UnitaryGaugeGroup w.r.t. trace dist: {times}.')