WIP: leakage-aware gauge optimization

pygsti/algorithms/gaugeopt.py

@@ -29,7 +29,7 @@ def gaugeopt_to_target(model, target_model, item_weights=None,
                        gauge_group=None, method='auto', maxiter=100000,
                        maxfev=None, tol=1e-8, oob_check_interval=0,
                        convert_model_to=None, return_all=False, comm=None,
-                       verbosity=0, check_jac=False):
+                       verbosity=0, check_jac=False, n_leak=0):
     """
     Optimize the gauge degrees of freedom of a model to that of a target.
 
@@ -170,7 +170,7 @@ def gaugeopt_to_target(model, target_model, item_weights=None,
     objective_fn, jacobian_fn = _create_objective_fn(
         model, target_model, item_weights,
         cptp_penalty_factor, spam_penalty_factor,
-        gates_metric, spam_metric, method, comm, check_jac)
+        gates_metric, spam_metric, method, comm, check_jac, n_leak)
 
     result = gaugeopt_custom(model, objective_fn, gauge_group, method,
                              maxiter, maxfev, tol, oob_check_interval,
@@ -303,9 +303,6 @@ def _call_jacobian_fn(gauge_group_el_vec):
 
     printer.log("--- Gauge Optimization (%s method, %s) ---" % (method, str(type(gauge_group))), 2)
     if method == 'ls':
-        #minSol  = _opt.least_squares(_call_objective_fn, x0, #jac=_call_jacobian_fn,
-        #                            max_nfev=maxfev, ftol=tol)
-        #solnX = minSol.x
         assert(_call_jacobian_fn is not None), "Cannot use 'ls' method unless jacobian is available"
         ralloc = _baseobjs.ResourceAllocation(comm)  # FUTURE: plumb up a resource alloc object?
         test_f = _call_objective_fn(x0)
@@ -353,7 +350,7 @@ def _call_jacobian_fn(gauge_group_el_vec):
 def _create_objective_fn(model, target_model, item_weights=None,
                          cptp_penalty_factor=0, spam_penalty_factor=0,
                          gates_metric="frobenius", spam_metric="frobenius",
-                         method=None, comm=None, check_jac=False):
+                         method=None, comm=None, check_jac=False, n_leak=0):
     """
     Creates the objective function and jacobian (if available)
     for gaugeopt_to_target
@@ -590,18 +587,48 @@ def _mock_objective_fn(v):
     else:
         # non-least-squares case where objective function returns a single float
         # and (currently) there's no analytic jacobian
+        dim = int(_np.sqrt(mxBasis.dim))
+        if n_leak > 0:
+            B = _tools.leading_dxd_submatrix_basis_vectors(dim - n_leak, dim, mxBasis)
+            """
+            ^ Need to do something else. 
+
+              In a leakage-friendly basis the operation matrices can be partitioned as 
+                  [comp     , semileak1]
+                  [semileak2, fulleak  ].
+              Instead of projecting only on to comp, we want to keep everything
+              except fullleak. ... But I don't think we can implement that just by
+              pre-multiplying by a projector.
+
+              I think this distinction might not be significant under certain idealized
+              assumptions, but small deviations from those conditions (which are certain
+              to happen in practice) might make it matter.
+            """
+            P = B @ B.T.conj()
+            if _np.linalg.norm(P.imag) > 1e-12:
+                raise ValueError()
+            else:
+                P = P.real
+            transform_mx_arg = (P, None)
+            # ^ The semantics of this tuple are defined by the frobeniusdist function
+            #   in the ExplicitOpModelCalc class. There are intended semantics for 
+            #   the second element of the tuple, but those aren't implemented yet so
+            #   for now I'm setting the second entry to None. -- Riley
+        else:
+            transform_mx_arg = (_np.eye(mxBasis.dim), None)
+
+        assert gates_metric != "frobeniustt"
+        assert spam_metric  != "frobeniustt"
+        # assert spam_metric == gates_metric
+        # ^ Erik and Corey said these are rarely used. I've removed support for
+        # them in this codepath (non-LS optimizer) in order to make it easier to
+        # read my updated code for leakage-aware metrics. It wouldn't be hard to
+        # add support back, but I just want to keep things simple. -- Riley
 
         def _objective_fn(gauge_group_el, oob_check):
             mdl = _transform_with_oob_check(model, gauge_group_el, oob_check)
             ret = 0
 
-            if gates_metric == "frobeniustt" or spam_metric == "frobeniustt":
-                full_target_model = target_model.copy()
-                full_target_model.convert_members_inplace("full")  # so we can gauge-transform the target model.
-                transformed_target = _transform_with_oob_check(full_target_model, gauge_group_el.inverse(), oob_check)
-            else:
-                transformed_target = None
-
             if cptp_penalty_factor > 0:
                 mdl.basis = mxBasis  # set basis for jamiolkowski iso
                 cpPenaltyVec = _cptp_penalty(mdl, cptp_penalty_factor, mdl.basis)
@@ -613,81 +640,97 @@ def _objective_fn(gauge_group_el, oob_check):
                 ret += _np.sum(spamPenaltyVec)
 
             if target_model is not None:
-                if gates_metric == "frobenius":
-                    if spam_metric == "frobenius":
-                        ret += mdl.frobeniusdist(target_model, None, 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
-                        ret += mdl.frobeniusdist(target_model, None, wts)
-
-                elif gates_metric == "frobeniustt":
-                    if spam_metric == "frobeniustt":
-                        ret += transformed_target.frobeniusdist(model, None, item_weights)
+                """
+                Leakage-aware metric supported, per implementation in mdl.frobeniusdist.
+                Refer to how mdl.frobeniusdist handles the case when transform_mx_arg
+                is a tuple in order to understand how the leakage-aware metric is defined.
+
+                  Idea: raise an error if mxBasis isn't leakage-friendly. 
+                      PROBLEM: if we're deep in the code then we end up needing to be
+                               working in a basis that has the identity matrix as its
+                               first element. The leakage-friendly basis doesn't even
+                               have the identity matrix as an element, let alone the first element
+
+                  TODO: dig into function calls below and see where we can
+                        access the mxBasis object. (I'm sure we can.)
+                        Looks like it isn't accessible within ExplicitOpModelCalc objects.
+                        It's most definitely available in ExplicitOpModel.       
+                """
+                if "frobenius" in gates_metric:
+                    if spam_metric == gates_metric:
+                        val = mdl.frobeniusdist(target_model, transform_mx_arg, item_weights)
                     else:
-                        wts = item_weights.copy(); wts['spam'] = 0.0
+                        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
-                        ret += transformed_target.frobeniusdist(model, None, 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)
+                    if "squared" in gates_metric:
+                        val = val ** 2
+                    ret += val
 
                 elif gates_metric == "fidelity":
+                    # Leakage-aware metric supported, using leaky_entanglement_fidelity.
                     for opLbl in mdl.operations:
                         wt = item_weights.get(opLbl, opWeight)
-                        ret += wt * (1.0 - _tools.entanglement_fidelity(
-                            target_model.operations[opLbl], mdl.operations[opLbl]))**2
+                        top = target_model.operations[opLbl].to_dense()
+                        mop = mdl.operations[opLbl].to_dense()
+                        ret += wt * (1.0 - _tools.leaky_entanglement_fidelity(top, mop, mxBasis, n_leak))**2
 
                 elif gates_metric == "tracedist":
+                    # Leakage-aware metric supported, using leaky_jtracedist.
                     for opLbl in mdl.operations:
                         wt = item_weights.get(opLbl, opWeight)
-                        ret += opWeight * _tools.jtracedist(
-                            target_model.operations[opLbl], mdl.operations[opLbl])
+                        top = target_model.operations[opLbl].to_dense()
+                        mop = mdl.operations[opLbl].to_dense()
+                        ret += wt * _tools.leaky_jtracedist(top, mop, mxBasis, n_leak)
 
                 else: raise ValueError("Invalid gates_metric: %s" % gates_metric)
 
-                if spam_metric == "frobenius":
-                    if gates_metric != "frobenius":  # otherwise handled above to match normalization in frobeniusdist
-                        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
-                        ret += mdl.frobeniusdist(target_model, None, wts)
-
-                elif spam_metric == "frobeniustt":
-                    if gates_metric != "frobeniustt":  # otherwise handled above to match normalization in frobeniusdist
-                        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
-                        ret += transformed_target.frobeniusdist(model, None, wts)
+                if "frobenius" in spam_metric and gates_metric == spam_metric:
+                    # We already handled SPAM error in this case. Just return.
+                    return ret
+
+                elif "frobenius" in spam_metric:
+                    # Leakage-aware metric supported in principle via implementation in
+                    # mdl.frobeniusdist (check implementation to see how it handles the
+                    # case when transform_mx_arg is a tuple).
+                    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)
+                    if "squared" in spam_metric:
+                        val = val ** 2
+                    ret += val 
 
                 elif spam_metric == "fidelity":
-                    for preplabel, prep in mdl.preps.items():
+                    # Leakage-aware metrics NOT available
+                    for preplabel, m_prep in mdl.preps.items():
                         wt = item_weights.get(preplabel, spamWeight)
-                        rhoMx1 = _tools.vec_to_stdmx(prep, mxBasis)
-                        rhoMx2 = _tools.vec_to_stdmx(
-                            target_model.preps[preplabel], mxBasis)
+                        rhoMx1 = _tools.vec_to_stdmx(m_prep.to_dense(), mxBasis)
+                        t_prep = target_model.preps[preplabel]
+                        rhoMx2 = _tools.vec_to_stdmx(t_prep.to_dense(), mxBasis)
                         ret += wt * (1.0 - _tools.fidelity(rhoMx1, rhoMx2))**2
-
-                    for povmlabel, povm in mdl.povms.items():
+ 
+                    for povmlabel in mdl.povms.keys():
                         wt = item_weights.get(povmlabel, spamWeight)
-                        ret += wt * (1.0 - _tools.povm_fidelity(
-                            mdl, target_model, povmlabel))**2
+                        fidelity = _tools.povm_fidelity(mdl, target_model, povmlabel)
+                        ret += wt * (1.0 - fidelity)**2
 
                 elif spam_metric == "tracedist":
-                    for preplabel, prep in mdl.preps.items():
+                    # Leakage-aware metrics NOT available.
+                    for preplabel, m_prep in mdl.preps.items():
                         wt = item_weights.get(preplabel, spamWeight)
-                        rhoMx1 = _tools.vec_to_stdmx(prep, mxBasis)
-                        rhoMx2 = _tools.vec_to_stdmx(
-                            target_model.preps[preplabel], mxBasis)
+                        rhoMx1 = _tools.vec_to_stdmx(m_prep.to_dense(), mxBasis)
+                        t_prep = target_model.preps[preplabel]
+                        rhoMx2 = _tools.vec_to_stdmx(t_prep.to_dense(), mxBasis)
                         ret += wt * _tools.tracedist(rhoMx1, rhoMx2)
 
-                    for povmlabel, povm in mdl.povms.items():
+                    for povmlabel in mdl.povms.keys():
                         wt = item_weights.get(povmlabel, spamWeight)
-                        ret += wt * (1.0 - _tools.povm_jtracedist(
-                            mdl, target_model, povmlabel))**2
+                        ret += wt * _tools.povm_jtracedist(mdl, target_model, povmlabel)
 
                 else: raise ValueError("Invalid spam_metric: %s" % spam_metric)
 

pygsti/baseobjs/basis.py

@@ -1760,3 +1760,76 @@ def create_simple_equivalent(self, builtin_basis_name=None):
                 if all([c.name == first_comp_name for c in self.component_bases]):
                     builtin_basis_name = first_comp_name  # if all components have the same name
         return BuiltinBasis(builtin_basis_name, self.elsize, sparse=self.sparse)
+
+
+class OuterProdBasis(TensorProdBasis):
+
+    def __init__(self, component_bases, name=None, longname=None, squeeze=True):
+        """
+        Let ⨂ denote an infix operator where (a ⨂ b) := numpy.tensordot(a, b, axes=0).
+
+        Suppose component_bases has length k. Under the default setting of squeeze=True,
+        the elements of this OuterProdBasis are all arrays of the form 
+
+            v1.squeeze() ⨂ v2.squeeze() ⨂ ... ⨂ vk.squeeze(),
+
+        where vi belongs to component_bases[i].elements. The definition is changed in the
+        natural way if squeeze=False.
+        """
+        TensorProdBasis.__init__(self, component_bases, name, longname)
+        if self.sparse:
+            raise NotImplementedError()
+        cnames = [c.name for c in self.component_bases]
+        if name is None:
+            self.name = "⨂".join(cnames)
+        if longname is None:
+            self.longname = "Outer-product basis with components " + ", ".join(cnames)
+        self._squeeze = squeeze
+        return
+
+    @property
+    def elshape(self):
+        shape = []
+        if self._squeeze:
+            for c in self.component_bases: 
+                shape.extend(d for d in c.elshape if d > 1)
+        else:
+            for c in self.component_bases:
+                shape.extend(c.elshape)
+        return tuple(shape)
+
+    def _lazy_build_elements(self):
+        compMxs = _np.zeros((self.size,) + self.elshape, 'complex')
+        comp_els = [c.elements for c in self.component_bases]
+        for i, factors in enumerate(_itertools.product(*comp_els)):
+            M = factors[0].squeeze() if self._squeeze else factors[0]
+            for f in factors[1:]:
+                if self._squeeze:
+                    f = f.squeeze()
+                M = _np.tensordot(M, f, axes=0)
+            compMxs[i] = M
+        self._elements = compMxs
+
+    def _lazy_build_labels(self):
+        self._labels = []
+        comp_lbls = [c.labels for c in self.component_bases]
+        for i, factor_lbls in enumerate(_itertools.product(*comp_lbls)):
+            self._labels.append('⨂'.join(factor_lbls))
+
+    def is_equivalent(self, other, sparseness_must_match=True):
+        if not sparseness_must_match:
+            raise NotImplementedError()
+        otherIsBasis = isinstance(other, OuterProdBasis)
+        if not otherIsBasis: return False  # can't be equal to a non-DirectSumBasis
+        return all([c1.is_equivalent(c2, True)
+                    for (c1, c2) in zip(self.component_bases, other.component_bases)])
+
+    def _copy_with_toggled_sparsity(self):
+        raise NotImplementedError()
+
+    def create_equivalent(self, builtin_basis_name):
+        raise NotImplementedError()
+
+    def create_simple_equivalent(self, builtin_basis_name=None):
+        raise NotImplementedError()
+

pygsti/circuits/circuit.py

@@ -801,12 +801,13 @@ def str(self, value):
         self._str = value
 
     def __hash__(self):
-        if not self._static:
-            _warnings.warn(("Editable circuit is being converted to read-only"
-                            " mode in order to hash it.  You should call"
-                            " circuit.done_editing() beforehand."))
-            self.done_editing()
-        return self._hash
+        return hash(self.tup)
+        # if not self._static:
+        #     _warnings.warn(("Editable circuit is being converted to read-only"
+        #                     " mode in order to hash it.  You should call"
+        #                     " circuit.done_editing() beforehand."))
+        #     self.done_editing()
+        # return self._hash
 
     def __len__(self):
         return len(self._labels)
@@ -978,10 +979,11 @@ def __eq__(self, x):
             if len(self) != len(x):
                 return False
             else:
-                if self._static and x._static:
-                    return self._hash == x._hash
-                else:
-                    return self.tup == x.tup
+                return self.tup == x.tup
+                # if self._static and x._static:
+                #     return self._hash == x._hash
+                # else:
+                #     return self.tup == x.tup
         elif x is None: 
             return False
         else:

pygsti/data/dataset.py

@@ -1176,7 +1176,7 @@ def _get_row(self, circuit):
         # needed because name-only Labels don't hash the same as strings
         # so key lookups need to be done at least with tuples of Labels.
         circuit = _cir.Circuit.cast(circuit)
-
+        circuit._static = False
         #Note: cirIndex value is either an int (non-static) or a slice (static)
         cirIndex = self.cirIndex[circuit]
         repData = self.repData[cirIndex] if (self.repData is not None) else None

pygsti/modelmembers/operations/fulltpop.py

@@ -48,6 +48,12 @@ class FullTPOp(_DenseOperator, _Torchable):
         of this state as a super-operator.  If None, certain functionality,
         such as access to Kraus operators, will be unavailable.
 
+        NOTE: this function might raise a weird error message if the passed
+              mx isn't in the pp or gm bases. I noticed this when using an
+              ExplicitBasis that is "leakage-friendly."
+              The implicit underlying assumption is that the first element
+              of this basis is the identity matrix.
+
     evotype : Evotype or str, optional
         The evolution type.  The special value `"default"` is equivalent
         to specifying the value of `pygsti.evotypes.Evotype.default_evotype`.