Ensure that gauge optimization incorporates instruments (#672)

This resolves #644 by making changes to the _create_objective_fn helper called by gaugeopt_to_target. The fix was easier than I expected for two reasons.

First, gauge optimization w.r.t. the Frobenius norm already incorporated
instruments by way of ExplicitOpModelCalc.frobeniusdist and
ExplicitOpModelCalc.residuals.

Second, gauge optimization w.r.t. trace
distance or infidelity can incorporate instruments just by working with
operation dicts from ``target_model._excalc()`` and ``mdl._excalc()``,
rather than the operation dicts in ``target_model`` and ``mdl``.

This PR's solution is somewhat suboptimal in that it increases reliance on ExplicitOpModelCalc, and a private method for constructing an ExplicitOpModelCalc at that. In a subsequent PR on general leakage modeling I'll revise _create_objective_fn to replace the current solution with something more elegant.

Author: rileyjmurray

pygsti/algorithms/gaugeopt.py

@@ -24,6 +24,8 @@
     TrivialGaugeGroupElement as _TrivialGaugeGroupElement,
     GaugeGroupElement as _GaugeGroupElement
 )
+from pygsti.models import ExplicitOpModel
+from pygsti.modelmembers.operations import LinearOperator as _LinearOperator
 
 from typing import Callable, Union, Optional
 
@@ -360,8 +362,10 @@ def _call_jacobian_fn(gauge_group_el_vec):
 
 GGElJacobian  = Union[None, Callable[[_GaugeGroupElement], _np.ndarray]]
 
+LabelLike = Union[str, _baseobjs.Label]
 
-def _create_objective_fn(model, target_model, item_weights: Optional[dict[str,float]]=None,
+
+def _create_objective_fn(model, target_model, item_weights: Optional[dict[LabelLike, float]]=None,
                          cptp_penalty_factor: float=0.0, spam_penalty_factor: float=0.0,
                          gates_metric="frobenius", spam_metric="frobenius",
                          method=None, comm=None, check_jac=False, n_leak=0) -> tuple[GGElObjective, GGElJacobian]:
@@ -625,8 +629,14 @@ def _mock_objective_fn(v):
             # ^ It would be equivalent to set this to a pair of IdentityOperator objects.
 
         def _objective_fn(gauge_group_el, oob_check):
-            mdl = _transform_with_oob_check(model, gauge_group_el, oob_check)
-            ret = 0
+            mdl : ExplicitOpModel = _transform_with_oob_check(model, gauge_group_el, oob_check)
+            mdl_ops : dict[_baseobjs.Label, _LinearOperator] = mdl._excalc().operations
+            tgt_ops : dict[_baseobjs.Label, _LinearOperator] = dict()
+            if target_model is not None:
+                tgt_ops.update(target_model._excalc().operations)
+            # ^ Use these dicts instead of mdl.operations and target_model.operations,
+            #   since these dicts are updated to include instruments.
+            ret = 0.0
 
             if cptp_penalty_factor > 0:
                 mdl.basis = mxBasis  # set basis for jamiolkowski iso
@@ -645,30 +655,27 @@ def _objective_fn(gauge_group_el, oob_check):
                 if spam_metric == gates_metric:
                     val = mdl.frobeniusdist(target_model, transform_mx_arg, item_weights)
                 else:
-                    wts = item_weights.copy()
-                    wts['spam'] = 0.0
-                    for k in wts:
-                        if k in mdl.preps or k in mdl.povms:
-                            wts[k] = 0.0
-                    val = mdl.frobeniusdist(target_model, transform_mx_arg, wts, n_leak)
+                    non_gates = {'spam'}.union(set(mdl.preps)).union(set(mdl.povms))
+                    temp_wts = {k: (0.0 if k in non_gates else v) for (k, v) in item_weights.items()}
+                    val = mdl.frobeniusdist(target_model, transform_mx_arg, temp_wts)
                 if "squared" in gates_metric:
                     val = val ** 2
                 ret += val
 
             elif gates_metric == "fidelity":
                 # If n_leak==0, then subspace_entanglement_fidelity is just entanglement_fidelity
-                for opLbl in mdl.operations:
+                for opLbl in mdl_ops:
                     wt = item_weights.get(opLbl, opWeight)
-                    top = target_model.operations[opLbl].to_dense()
-                    mop = mdl.operations[opLbl].to_dense()
+                    top = tgt_ops[opLbl].to_dense()
+                    mop = mdl_ops[opLbl].to_dense()
                     ret += wt * (1.0 - _tools.subspace_entanglement_fidelity(top, mop, mxBasis, n_leak))**2
 
             elif gates_metric == "tracedist":
                 # If n_leak==0, then subspace_jtracedist is just jtracedist.
-                for opLbl in mdl.operations:
+                for opLbl in mdl_ops:
                     wt = item_weights.get(opLbl, opWeight)
-                    top = target_model.operations[opLbl].to_dense()
-                    mop = mdl.operations[opLbl].to_dense()
+                    top = tgt_ops[opLbl].to_dense()
+                    mop = mdl_ops[opLbl].to_dense()
                     ret += wt * _tools.subspace_jtracedist(top, mop, mxBasis, n_leak)
 
             else:
@@ -680,11 +687,9 @@ def _objective_fn(gauge_group_el, oob_check):
 
             if "frobenius" in spam_metric:
                 # SPAM and gates can have different choices for squared vs non-squared.
-                wts = item_weights.copy(); wts['gates'] = 0.0
-                for k in wts:
-                    if k in mdl.operations or k in mdl.instruments:
-                        wts[k] = 0.0
-                val = mdl.frobeniusdist(target_model, transform_mx_arg, wts)
+                non_spam = {'gates'}.union(set(mdl_ops)) # use mdl_ops, not mdl.operations!
+                temp_wts = {k: (0.0 if k in non_spam else v) for (k, v) in item_weights.items()}
+                val = mdl.frobeniusdist(target_model, transform_mx_arg, temp_wts)
                 if "squared" in spam_metric:
                     val = val ** 2
                 ret += val 

pygsti/modelmembers/operations/linearop.py

@@ -506,9 +506,8 @@ def jtracedist(self, other_op, transform=None, inv_transform=None):
         if transform is None and inv_transform is None:
             return _ot.jtracedist(self.to_dense("minimal"), other_op.to_dense("minimal"))
         else:
-            return _ot.jtracedist(_np.dot(
-                inv_transform, _np.dot(self.to_dense("minimal"), transform)),
-                other_op.to_dense("minimal"))
+            arg = inv_transform @ self.to_dense("minimal") @ transform
+            return _ot.jtracedist(arg, other_op.to_dense("minimal"))
 
     def diamonddist(self, other_op, transform=None, inv_transform=None):
         """
@@ -535,9 +534,8 @@ def diamonddist(self, other_op, transform=None, inv_transform=None):
         if transform is None and inv_transform is None:
             return _ot.diamonddist(self.to_dense("minimal"), other_op.to_dense("minimal"))
         else:
-            return _ot.diamonddist(_np.dot(
-                inv_transform, _np.dot(self.to_dense("minimal"), transform)),
-                other_op.to_dense("minimal"))
+            arg = inv_transform @ self.to_dense("minimal") @ transform
+            return _ot.diamonddist(arg, other_op.to_dense("minimal"))
 
     def transform_inplace(self, s):
         """
@@ -563,7 +561,7 @@ def transform_inplace(self, s):
         """
         Smx = s.transform_matrix
         Si = s.transform_matrix_inverse
-        self.set_dense(_np.dot(Si, _np.dot(self.to_dense("minimal"), Smx)))
+        self.set_dense(Si @ self.to_dense("minimal") @ Smx)
 
     def spam_transform_inplace(self, s, typ):
         """
@@ -591,9 +589,9 @@ def spam_transform_inplace(self, s, typ):
         None
         """
         if typ == 'prep':
-            self.set_dense(_np.dot(s.transform_matrix_inverse, self.to_dense("minimal")))
+            self.set_dense(s.transform_matrix_inverse @ self.to_dense("minimal"))
         elif typ == 'effect':
-            self.set_dense(_np.dot(self.to_dense("minimal"), s.transform_matrix))
+            self.set_dense(self.to_dense("minimal") @ s.transform_matrix)
         else:
             raise ValueError("Invalid `typ` argument: %s" % typ)
 
@@ -627,7 +625,7 @@ def depolarize(self, amount):
         else:
             assert(len(amount) == self.dim - 1)
             D = _np.diag([1] + list(1.0 - _np.array(amount, 'd')))
-        self.set_dense(_np.dot(D, self.to_dense("minimal")))
+        self.set_dense(D @  self.to_dense("minimal"))
 
     def rotate(self, amount, mx_basis="gm"):
         """
@@ -658,7 +656,7 @@ def rotate(self, amount, mx_basis="gm"):
         None
         """
         rotnMx = _ot.rotation_gate_mx(amount, mx_basis)
-        self.set_dense(_np.dot(rotnMx, self.to_dense("minimal")))
+        self.set_dense(rotnMx @ self.to_dense("minimal"))
 
     def deriv_wrt_params(self, wrt_filter=None):
         """

pygsti/models/explicitcalc.py

@@ -23,7 +23,24 @@
 
 from typing import Union
 
-TransformMxPair = tuple[_np.ndarray, Union[_np.ndarray, _mt.IdentityOperator]]
+
+TransformMxPair = tuple[
+    Union[_np.ndarray, _mt.IdentityOperator],
+    Union[_np.ndarray, _mt.IdentityOperator]
+]
+
+
+def to_transform_mx_pair(tmx_arg: Union[None, _np.ndarray, TransformMxPair]) -> TransformMxPair:
+    if tmx_arg is None:
+        P = invP = _mt.IdentityOperator()
+    elif isinstance(tmx_arg, _np.ndarray):
+        P = tmx_arg
+        invP = _np.linalg.pinv(P)
+    else:
+        P, invP = tmx_arg
+        assert _mt.is_operatorlike(P)
+        assert _mt.is_operatorlike(invP)
+    return P, invP
 
 # Tolerace for matrix_rank when finding rank of a *normalized* projection
 # matrix.  This is a legitimate tolerace since a normalized projection matrix
@@ -156,16 +173,7 @@ def frobeniusdist(self, other_calc, transform_mx: Union[None, _np.ndarray, Trans
         -------
         float
         """
-        if transform_mx is None:
-            P, invP = None, None
-            # ^ It would be equivalent to use P = invP = _mt.IdentityOperator()
-        elif isinstance(transform_mx, _np.ndarray):
-            P = transform_mx
-            invP = _np.linalg.pinv(P)
-        else:
-            P, invP = transform_mx
-            assert _mt.is_operatorlike(P)
-            assert _mt.is_operatorlike(invP)
+        P, invP = to_transform_mx_pair(transform_mx)
 
         d = 0.0
         nSummands = 0.0
@@ -231,14 +239,12 @@ def residuals(self, other_calc, transform_mx=None, item_weights=None):
             The (weighted) number of elements accounted for by the residuals.
         """
         resids = []
-        T = transform_mx
         nSummands = 0.0
         if item_weights is None: item_weights = {}
         sqrt_itemWeights = {k: _np.sqrt(v) for k, v in item_weights.items()}
         opWeight = sqrt_itemWeights.get('gates', 1.0)
         spamWeight = sqrt_itemWeights.get('spam', 1.0)
-        Ti = None if T is None else _np.linalg.pinv(T)
-        # ^ TODO: generalize inverse op (call T.inverse() if T were a "transform" object?)
+        T, Ti = to_transform_mx_pair(transform_mx)
 
         for opLabel, gate in self.operations.items():
             wt = sqrt_itemWeights.get(opLabel, opWeight)
@@ -293,40 +299,23 @@ def jtracedist(self, other_calc, transform_mx=None, include_spam=True):
         -------
         float
         """
-        T = transform_mx
+        T, Ti = to_transform_mx_pair(transform_mx)
         d = 0  # spam difference
         nSummands = 0  # for spam terms
 
-        if T is not None:
-            Ti = _np.linalg.inv(T)
-            dists = [gate.jtracedist(other_calc.operations[lbl], T, Ti)
-                     for lbl, gate in self.operations.items()]
-
-            #Just use frobenius distance between spam vecs, since jtracedist
-            # doesn't really make sense
-            if include_spam:
-                for lbl, rhoV in self.preps.items():
-                    d += rhoV.frobeniusdist_squared(other_calc.preps[lbl], T, Ti)
-                    nSummands += rhoV.dim
+        dists = [gate.jtracedist(other_calc.operations[lbl], T, Ti)
+                    for lbl, gate in self.operations.items()]
 
-                for lbl, Evec in self.effects.items():
-                    d += Evec.frobeniusdist_squared(other_calc.effects[lbl], T, Ti)
-                    nSummands += Evec.dim
+        # Just use frobenius distance between spam vecs, since jtracedist
+        # doesn't really make sense
+        if include_spam:
+            for lbl, rhoV in self.preps.items():
+                d += rhoV.frobeniusdist_squared(other_calc.preps[lbl], T, Ti)
+                nSummands += rhoV.dim
 
-        else:
-            dists = [gate.jtracedist(other_calc.operations[lbl])
-                     for lbl, gate in self.operations.items()]
-
-            #Just use frobenius distance between spam vecs, since jtracedist
-            # doesn't really make sense
-            if include_spam:
-                for lbl, rhoV in self.preps.items():
-                    d += rhoV.frobeniusdist_squared(other_calc.preps[lbl])
-                    nSummands += rhoV.dim
-
-                for lbl, Evec in self.effects.items():
-                    d += Evec.frobeniusdist_squared(other_calc.effects[lbl])
-                    nSummands += Evec.dim
+            for lbl, Evec in self.effects.items():
+                d += Evec.frobeniusdist_squared(other_calc.effects[lbl], T, Ti)
+                nSummands += Evec.dim
 
         spamVal = _np.sqrt(d / nSummands) if (nSummands > 0) else 0
         return max(dists) + spamVal
@@ -358,41 +347,24 @@ def diamonddist(self, other_calc, transform_mx=None, include_spam=True):
         -------
         float
         """
-        T = transform_mx
+        T, Ti = to_transform_mx_pair(transform_mx)
         d = 0  # spam difference
         nSummands = 0  # for spam terms
 
-        if T is not None:
-            Ti = _np.linalg.inv(T)
-            dists = [gate.diamonddist(other_calc.operations[lbl], T, Ti)
-                     for lbl, gate in self.operations.items()]
+        dists = [gate.diamonddist(other_calc.operations[lbl], T, Ti)
+                    for lbl, gate in self.operations.items()]
 
-            #Just use frobenius distance between spam vecs, since jtracedist
-            # doesn't really make sense
-            if include_spam:
-                for lbl, rhoV in self.preps.items():
-                    d += rhoV.frobeniusdist_squared(other_calc.preps[lbl], T, Ti)
-                    nSummands += rhoV.dim
+        # Just use frobenius distance between spam vecs, since diamonddist
+        # doesn't really make sense
+        if include_spam:
+            for lbl, rhoV in self.preps.items():
+                d += rhoV.frobeniusdist_squared(other_calc.preps[lbl], T, Ti)
+                nSummands += rhoV.dim
 
-                for lbl, Evec in self.effects.items():
+            for lbl, Evec in self.effects.items():
                     d += Evec.frobeniusdist_squared(other_calc.effects[lbl], T, Ti)
                     nSummands += Evec.dim
 
-        else:
-            dists = [gate.diamonddist(other_calc.operations[lbl])
-                     for lbl, gate in self.operations.items()]
-
-            #Just use frobenius distance between spam vecs, since jtracedist
-            # doesn't really make sense
-            if include_spam:
-                for lbl, rhoV in self.preps.items():
-                    d += rhoV.frobeniusdist_squared(other_calc.preps[lbl])
-                    nSummands += rhoV.dim
-
-                for lbl, Evec in self.effects.items():
-                    d += Evec.frobeniusdist_squared(other_calc.effects[lbl])
-                    nSummands += Evec.dim
-
         spamVal = _np.sqrt(d / nSummands) if (nSummands > 0) else 0
         return max(dists) + spamVal
 
@@ -420,7 +392,7 @@ def deriv_wrt_params(self):
             eo += obj.hilbert_schmidt_size
 
         if self.interposer is not None:
-            deriv = _np.dot(deriv, self.interposer.deriv_op_params_wrt_model_params())
+            deriv = deriv @ self.interposer.deriv_op_params_wrt_model_params()
 
         return deriv
 
@@ -499,14 +471,13 @@ def _buildup_dpg(self):
             for j in range(dim):  # *generator* mx, not gauge mx itself
                 unitMx = _bc.mut(i, j, dim)
                 for lbl, rhoVec in self_preps.items():
-                    mdlDeriv_preps[lbl] = _np.dot(unitMx, rhoVec)
+                    mdlDeriv_preps[lbl] = unitMx @ rhoVec
                 for lbl, EVec in self_effects.items():
-                    mdlDeriv_effects[lbl] = -_np.dot(EVec.T, unitMx).T
+                    mdlDeriv_effects[lbl] = -(EVec.T @ unitMx).T
 
                 for lbl, gate in self_operations.items():
                     #if isinstance(gate,_op.DenseOperator):
-                    mdlDeriv_ops[lbl] = _np.dot(unitMx, gate) - \
-                        _np.dot(gate, unitMx)
+                    mdlDeriv_ops[lbl] = (unitMx @ gate) - (gate @ unitMx)
                     #else:
                     #    #use acton... maybe throw error if dim is too large (maybe above?)
                     #    deriv = _np.zeros((dim,dim),'d')
@@ -566,7 +537,7 @@ def nongauge_and_gauge_spaces(self, item_weights=None, non_gauge_mix_mx=None):
             #for each column of gen_dG, which is a gauge direction in model parameter space,
             # we add some amount of non-gauge direction, given by coefficients of the
             # numNonGaugeParams non-gauge directions.
-            orthog_to = gauge_space + _np.dot(nonGaugeDirections, non_gauge_mix_mx)  # add non-gauge components in
+            orthog_to = gauge_space + nonGaugeDirections @ non_gauge_mix_mx  # add non-gauge components in
             # dims: (nParams,n_gauge_params) + (nParams,n_non_gauge_params) * (n_non_gauge_params,n_gauge_params)
             # non_gauge_mix_mx is a (n_non_gauge_params,n_gauge_params) matrix whose i-th column specifies the
             #  coefficients to multiply each of the non-gauge directions by before adding them to the i-th
@@ -583,7 +554,7 @@ def nongauge_and_gauge_spaces(self, item_weights=None, non_gauge_mix_mx=None):
                 metric_diag[vec.gpindices] = item_weights.get(lbl, spamWeight)
             metric = _np.diag(metric_diag)
             #OLD: gen_ndG = _mt.nullspace(_np.dot(gen_dG.T,metric))
-            orthog_to = _np.dot(metric.T, gauge_space)
+            orthog_to = metric.T @ gauge_space
 
         else:
             orthog_to = gauge_space
@@ -663,9 +634,9 @@ def _gauge_orbit_curvature(self, item_weights=None, non_gauge_mix_mx=None):
                         unitMx_j = _bc.mut(j1, j2, dim)
                         antiComm = (unitMx_i @ unitMx_j + unitMx_j @ unitMx_i)
                         for lbl, rhoVec in self_preps.items():
-                            mdlHess_preps[lbl] = 0.5 * _np.dot(antiComm, rhoVec)
+                            mdlHess_preps[lbl] = 0.5 * (antiComm @ rhoVec)
                         for lbl, EVec in self_effects.items():
-                            mdlHess_effects[lbl] = 0.5 * _np.dot(EVec.T, antiComm).T
+                            mdlHess_effects[lbl] = 0.5 * (EVec.T @ antiComm).T
                         for lbl, gate in self_operations.items():
                             mdlHess_ops[lbl] = 0.5 * (antiComm @ gate + gate @ antiComm) \
                                 - unitMx_i @ gate @ unitMx_j - unitMx_j @ gate @ unitMx_i
@@ -785,8 +756,8 @@ def nongauge_projector(self, item_weights=None, non_gauge_mix_mx=None):
 
         # ORIG WAY: use pseudo-inverse to normalize projector.  Ran into problems where
         #  default rcond == 1e-15 didn't work for 2-qubit case, but still more stable than inv method below
-        P = _np.dot(gen_ndG, _np.transpose(gen_ndG))  # almost a projector, but cols of dG are not orthonormal
-        Pp = _np.dot(_np.linalg.pinv(P, rcond=1e-7), P)  # make P into a true projector (onto gauge space)
+        P = gen_ndG @ gen_ndG.T  # almost a projector, but cols of dG are not orthonormal
+        Pp = _np.linalg.pinv(P, rcond=1e-7) @ P  # make P into a true projector (onto gauge space)
 
         # ALT WAY: use inverse of dG^T*dG to normalize projector (see wikipedia on projectors, dG => A)
         #  This *should* give the same thing as above, but numerical differences indicate the pinv method