Ensure that gauge optimization incorporates instruments

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