Leakage-aware GST

jupyter_notebooks/Examples/Leakage-automagic.ipynb

@@ -0,0 +1,87 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pygsti.modelpacks import smq1Q_XY, smq1Q_ZN\n",
+    "from pygsti.tools.leakage import leaky_qubit_model_from_pspec, construct_leakage_report\n",
+    "from pygsti.data import simulate_data\n",
+    "from pygsti.protocols import StandardGST, ProtocolData"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## GST: modeling a leaky qubit as a qutrit\n",
+    "\n",
+    "This short notebook shows how (data from) an experiment design for a two-level system can be used to fit a three-level sytem model, and how to generate a special report to provide insights for these models. The report includes special gate error metrics that reflect the distinguished role of the first two levels in the three-level system."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mp = smq1Q_XY\n",
+    "ed = mp.create_gst_experiment_design(max_max_length=32)\n",
+    "tm3 = leaky_qubit_model_from_pspec(mp.processor_spec(), mx_basis='l2p1')\n",
+    "# ^ We could use basis = 'gm' instead of 'l2p1'. We prefer 'l2p1'\n",
+    "#   because it makes process matrices easier to interpret in leakage\n",
+    "#   modeling.\n",
+    "ds = simulate_data(tm3, ed.all_circuits_needing_data, num_samples=1000, seed=1997)\n",
+    "gst = StandardGST( modes=('CPTPLND',), target_model=tm3, verbosity=2)\n",
+    "pd = ProtocolData(ed, ds)\n",
+    "res = gst.run(pd)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "report_dir = 'example_files/leakage-report-automagic'\n",
+    "report_object, updated_res = construct_leakage_report(res, title='easy leakage analysis!')\n",
+    "# ^ Each estimate in updated_res has a new gauge-optimized model.\n",
+    "#   The gauge optimization was done to reflect how our target gates\n",
+    "#   are only _really_ defined on the first two levels of our\n",
+    "#   three-level system.\n",
+    "#   \n",
+    "report_object.write_html(report_dir)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "leak311",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

pygsti/algorithms/gaugeopt.py

@@ -20,7 +20,12 @@
 from pygsti import tools as _tools
 from pygsti.tools import mpitools as _mpit
 from pygsti.tools import slicetools as _slct
-from pygsti.models.gaugegroup import TrivialGaugeGroupElement as _TrivialGaugeGroupElement
+from pygsti.models.gaugegroup import (
+    TrivialGaugeGroupElement as _TrivialGaugeGroupElement,
+    GaugeGroupElement as _GaugeGroupElement
+)
+
+from typing import Callable, Union, Optional
 
 
 def gaugeopt_to_target(model, target_model, item_weights=None,
@@ -29,7 +34,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 +175,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,
@@ -307,9 +312,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)
@@ -354,10 +356,15 @@ def _call_jacobian_fn(gauge_group_el_vec):
         return newModel
 
 
-def _create_objective_fn(model, target_model, item_weights=None,
-                         cptp_penalty_factor=0, spam_penalty_factor=0,
+GGElObjective = Callable[[_GaugeGroupElement,bool], Union[float, _np.ndarray]]
+
+GGElJacobian  = Union[None, Callable[[_GaugeGroupElement], _np.ndarray]]
+
+
+def _create_objective_fn(model, target_model, item_weights: Optional[dict[str,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):
+                         method=None, comm=None, check_jac=False, n_leak=0) -> tuple[GGElObjective, GGElJacobian]:
     """
     Creates the objective function and jacobian (if available)
     for gaugeopt_to_target
@@ -595,17 +602,32 @@ def _mock_objective_fn(v):
         # non-least-squares case where objective function returns a single float
         # and (currently) there's no analytic jacobian
 
+        assert gates_metric != "frobeniustt"
+        assert spam_metric  != "frobeniustt"
+        # ^ PR #410 removed support for Frobenius transform-target metrics in this codepath.
+
+        dim = int(_np.sqrt(mxBasis.dim))
+        if n_leak > 0:
+            B = _tools.leading_dxd_submatrix_basis_vectors(dim - n_leak, dim, mxBasis)
+            P = B @ B.T.conj()
+            if _np.linalg.norm(P.imag) > 1e-12:
+                msg  = f"Attempting to run leakage-aware gauge optimization with basis {mxBasis}\n"
+                msg +=  "is resulting an orthogonal projector onto the computational subspace that\n"
+                msg +=  "is not real-valued. Try again with a different basis, like 'l2p1' or 'gm'."
+                raise ValueError(msg)
+            else:
+                P = P.real
+            transform_mx_arg = (P, _tools.matrixtools.IdentityOperator())
+            # ^ The semantics of this tuple are defined by the frobeniusdist function
+            #   in the ExplicitOpModelCalc class.
+        else:
+            transform_mx_arg = None
+            # ^ 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
 
-            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)
@@ -616,84 +638,86 @@ def _objective_fn(gauge_group_el, oob_check):
                 spamPenaltyVec = _spam_penalty(mdl, spam_penalty_factor, mdl.basis)
                 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)
-                    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 += transformed_target.frobeniusdist(model, None, wts)
-
-                elif gates_metric == "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
-
-                elif gates_metric == "tracedist":
-                    for opLbl in mdl.operations:
-                        wt = item_weights.get(opLbl, opWeight)
-                        ret += opWeight * _tools.jtracedist(
-                            target_model.operations[opLbl], mdl.operations[opLbl])
-
-                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)
-
-                elif spam_metric == "fidelity":
-                    for preplabel, 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)
-                        ret += wt * (1.0 - _tools.fidelity(rhoMx1, rhoMx2))**2
-
-                    for povmlabel, povm in mdl.povms.items():
-                        wt = item_weights.get(povmlabel, spamWeight)
-                        ret += wt * (1.0 - _tools.povm_fidelity(
-                            mdl, target_model, povmlabel))**2
-
-                elif spam_metric == "tracedist":
-                    for preplabel, 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)
-                        ret += wt * _tools.tracedist(rhoMx1, rhoMx2)
-
-                    for povmlabel, povm in mdl.povms.items():
-                        wt = item_weights.get(povmlabel, spamWeight)
-                        ret += wt * (1.0 - _tools.povm_jtracedist(
-                            mdl, target_model, povmlabel))**2
-
-                else: raise ValueError("Invalid spam_metric: %s" % spam_metric)
+            if target_model is None:
+                return ret
+
+            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
+                    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)
+                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:
+                    wt = item_weights.get(opLbl, opWeight)
+                    top = target_model.operations[opLbl].to_dense()
+                    mop = mdl.operations[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:
+                    wt = item_weights.get(opLbl, opWeight)
+                    top = target_model.operations[opLbl].to_dense()
+                    mop = mdl.operations[opLbl].to_dense()
+                    ret += wt * _tools.subspace_jtracedist(top, mop, mxBasis, n_leak)
+
+            else:
+                raise ValueError("Invalid gates_metric: %s" % gates_metric)
+
+            if "frobenius" in spam_metric and gates_metric == spam_metric:
+                # We already handled SPAM error in this case. Just return.
+                return ret
+
+            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)
+                if "squared" in spam_metric:
+                    val = val ** 2
+                ret += val 
+
+            elif spam_metric == "fidelity":
+                # Leakage-aware metrics NOT available
+                for preplabel, m_prep in mdl.preps.items():
+                    wt = item_weights.get(preplabel, spamWeight)
+                    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 in mdl.povms.keys():
+                    wt = item_weights.get(povmlabel, spamWeight)
+                    fidelity = _tools.povm_fidelity(mdl, target_model, povmlabel)
+                    ret += wt * (1.0 - fidelity)**2
+
+            elif spam_metric == "tracedist":
+                # Leakage-aware metrics NOT available.
+                for preplabel, m_prep in mdl.preps.items():
+                    wt = item_weights.get(preplabel, spamWeight)
+                    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 in mdl.povms.keys():
+                    wt = item_weights.get(povmlabel, spamWeight)
+                    ret += wt * _tools.povm_jtracedist(mdl, target_model, povmlabel)
+
+            else:
+                raise ValueError("Invalid spam_metric: %s" % spam_metric)
 
             return ret
 

pygsti/baseobjs/basis.py

@@ -794,6 +794,22 @@ def ellookup(self):
                 self._lazy_build_labels()
             self._ellookup = {lbl: el for lbl, el in zip(self._labels, self._elements)}
         return self._ellookup
+
+    @property
+    def elindlookup(self) -> dict:
+        """
+        A dictionary mapping labels to the index in self.elements that holds
+        that label's basis element.
+
+        Returns
+        -------
+        dict
+        """
+        if self._ellookup is None:
+            if self._labels is None:
+                self._lazy_build_labels()
+            self._elindlookup = {lbl: ind for ind, lbl in enumerate(self._labels)}
+        return self._elindlookup
 
     @property
     def elements(self):