Feature ddist wildcard integration

pygsti/objectivefns/wildcardbudget.py

@@ -539,7 +539,7 @@ def __init__(self, primitive_op_labels, start_budget=0.0, idle_name=None):
         nParams = len(set(self.primOpLookup.values()))
         if isinstance(start_budget, dict):
             Wvec = _np.zeros(nParams, 'd')
-            for op, val in start_budget.items:
+            for op, val in start_budget.items():
                 Wvec[self.primOpLookup[op]] = val
         else:
             Wvec = _np.array([start_budget] * nParams)
@@ -914,3 +914,122 @@ def update_circuit_probs(probs, freqs, circuit_budget):
     #assert(_np.isclose(W, compTVD)), "TVD mismatch!"
 
     return updated_qvec
+
+class PrimitiveOpsDiamondDistanceWildcardBudget(PrimitiveOpsWildcardBudget):
+    """
+    A wildcard budget containing one parameter per "primitive operation"
+    based on a single parameter model wherein each of the wildcard parameters
+    is set proportionally to the value of a gate's gauge optimized diamond distance.
+
+    This class extends PrimitiveOpsWildcardBudget to also internally store
+    the value of the alpha parameter (how much the diamond distance needed to
+    be scaled to restore consistency and to store the values of the diamond
+    distances used to derive the wildcard values.
+
+    A parameter's absolute value gives the amount of "slack", or
+    "wildcard budget" that is allocated per that particular primitive
+    operation.
+
+    Primitive operations are the components of circuit layers, and so
+    the wilcard budget for a circuit is just the sum of the (abs vals of)
+    the parameters corresponding to each primitive operation in the circuit.
+
+    Parameters
+    ----------
+    primitive_op_labels : iterable or dict
+        A list of primitive-operation labels, e.g. `Label('Gx',(0,))`,
+        which give all the possible primitive ops (components of circuit
+        layers) that will appear in circuits.  Each one of these operations
+        will be assigned it's own independent element in the wilcard-vector.
+        A dictionary can be given whose keys are Labels and whose values are
+        0-based parameter indices.  In the non-dictionary case, each label gets
+        it's own parameter.  Dictionaries allow multiple labels to be associated
+        with the *same* wildcard budget parameter,
+        e.g. `{Label('Gx',(0,)): 0, Label('Gy',(0,)): 0}`.
+        If `'SPAM'` is included as a primitive op, this value correspond to a
+        uniform "SPAM budget" added to each circuit.
+
+    start_budget : float or dict, optional
+        An initial value to set all the parameters to (if a float), or a
+        dictionary mapping primitive operation labels to initial values.
+
+    ddists : dict, optional
+        A dictionary of diamond distance values, should be in the same format as
+        start_budget.
+
+    alpha : float, optional
+        The value of the alpha parameter that gives the diamond distance scaling
+        for this wildcard model.
+
+    idle_name : str, optional
+            The gate name to be used for the 1-qubit idle gate.  If not `None`, then
+            circuit budgets are computed by considering layers of the circuit as being
+            "padded" with `1-qubit` idles gates on any empty lines.
+    """
+
+    def __init__(self, primitive_op_labels, start_budget=0.0, ddists=None, alpha=0, idle_name=None):
+        """
+        Create a new PrimitiveOpsWildcardBudget.
+
+        Parameters
+        ----------
+        primitive_op_labels : iterable or dict
+            A list of primitive-operation labels, e.g. `Label('Gx',(0,))`,
+            which give all the possible primitive ops (components of circuit
+            layers) that will appear in circuits.  Each one of these operations
+            will be assigned it's own independent element in the wilcard-vector.
+            A dictionary can be given whose keys are Labels and whose values are
+            0-based parameter indices.  In the non-dictionary case, each label gets
+            it's own parameter.  Dictionaries allow multiple labels to be associated
+            with the *same* wildcard budget parameter,
+            e.g. `{Label('Gx',(0,)): 0, Label('Gy',(0,)): 0}`.
+            If `'SPAM'` is included as a primitive op, this value correspond to a
+            uniform "SPAM budget" added to each circuit.
+
+        start_budget : float or dict, optional
+            An initial value to set all the parameters to (if a float), or a
+            dictionary mapping primitive operation labels to initial values.
+
+        idle_name : str, optional
+            The gate name to be used for the 1-qubit idle gate.  If not `None`, then
+            circuit budgets are computed by considering layers of the circuit as being
+            "padded" with `1-qubit` idles gates on any empty lines.
+
+        """
+
+        if ddists is not None:
+            if isinstance(ddists, dict):
+                self.ddists= list(ddists.values())
+            else:
+                self.ddists= ddists
+        self.alpha= alpha
+
+        super().__init__(primitive_op_labels, start_budget, idle_name)
+
+
+    @property
+    def description(self):
+        """
+        A dictionary of quantities describing this budget.
+
+        Return the contents of this budget in a dictionary containing
+        (description, value) pairs for each element name.
+
+        Returns
+        -------
+        dict
+            Keys are primitive op labels and values are (description_string, value) tuples.
+        """
+        wildcardDict = {}
+        for lbl, index in self.primOpLookup.items():
+            if lbl == "SPAM": continue  # treated separately below
+            wildcardDict[lbl] = ('budget per each instance %s' % str(lbl), pos(self.wildcard_vector[index]))
+        if self.spam_index is not None:
+            wildcardDict['SPAM'] = ('uniform per-circuit SPAM budget', pos(self.wildcard_vector[self.spam_index]))
+
+        #Add an entry for the alpha parameter.
+        wildcardDict['alpha']= ('Fraction of diamond distance needed as wildcard error', self.alpha)
+
+        return wildcardDict
+
+

pygsti/protocols/gst.py

@@ -43,6 +43,9 @@
 from pygsti.tools.legacytools import deprecate as _deprecated_fn
 
 
+
+
+
 #For results object:
 ROBUST_SUFFIX_LIST = [".robust", ".Robust", ".robust+", ".Robust+"]
 DEFAULT_BAD_FIT_THRESHOLD = 2.0
@@ -524,11 +527,28 @@ class GSTBadFitOptions(_NicelySerializable):
         GST fit is considered satisfactory (and no "bad-fit" processing is needed).
 
     actions : tuple, optional
-        Actions to take when a GST fit is unsatisfactory.
+        Actions to take when a GST fit is unsatisfactory. Allowed actions include:
+        - 'wildcard': Find an admissable wildcard model...
+        - 'ddist_wildcard': Fits a single parameter wildcard model in which
+          the amount of wildcard error added to an operation is proportional
+          to the diamond distance between that operation and the target.
+        - 'Robust+': ...
+        - 'robust+': ...
+        - 'robust': ...
+        - 'do nothing': ...
 
     wildcard_budget_includes_spam : bool, optional
         Include a SPAM budget within the wildcard budget used to process
         the `"wildcard"` action.
+
+    wildcard_L1_weights :
+
+    wildcard_primitive_op_labels:
+
+    wildcard_methods: tuple, optional (default ('neldermead',))
+
+    wildcard_inadmissable_action: str, optional (default 'print')
+
     """
 
     @classmethod
@@ -556,7 +576,7 @@ def __init__(self, threshold=DEFAULT_BAD_FIT_THRESHOLD, actions=(),
                  wildcard_L1_weights=None, wildcard_primitive_op_labels=None,
                  wildcard_initial_budget=None, wildcard_methods=('neldermead',),
                  wildcard_inadmissable_action='print'):
-        valid_actions = ('wildcard', 'Robust+', 'Robust', 'robust+', 'robust', 'do nothing')
+        valid_actions = ('wildcard', 'ddist_wildcard', 'Robust+', 'Robust', 'robust+', 'robust', 'do nothing')
         if not all([(action in valid_actions) for action in actions]):
             raise ValueError("Invalid action in %s! Allowed actions are %s" % (str(actions), str(valid_actions)))
         self.threshold = float(threshold)
@@ -1834,6 +1854,31 @@ def _add_badfit_estimates(results, base_estimate_label, badfit_options,
             #    printer.warning("Failed to get wildcard budget - continuing anyway.  Error was:\n" + str(e))
             #    new_params['unmodeled_error'] = None
             continue  # no need to add a new estimate - we just update the base estimate
+
+        elif badfit_typ == 'ddist_wildcard':
+
+            #If this estimate is the target model then skip adding the diamond distance wildcard.
+            if base_estimate_label != 'Target':
+                try:
+                    budget = _compute_wildcard_budget_ddist_model(base_estimate,objfn_cache, mdc_objfn, parameters, badfit_options, printer - 1)
+
+                    base_estimate.extra_parameters['ddist_wildcard' + "_unmodeled_error"] = budget
+                    base_estimate.extra_parameters['ddist_wildcard' + "_unmodeled_active_constraints"] \
+                        = None
+
+                    base_estimate.extra_parameters["unmodeled_error"] = budget
+                    base_estimate.extra_parameters["unmodeled_active_constraints"] = None
+                except NotImplementedError as e:
+                    printer.warning("Failed to get wildcard budget - continuing anyway.  Error was:\n" + str(e))
+                    new_params['unmodeled_error'] = None
+                #except AssertionError as e:
+                #    printer.warning("Failed to get wildcard budget - continuing anyway.  Error was:\n" + str(e))
+                #    new_params['unmodeled_error'] = None
+                continue  # no need to add a new estimate - we just update the base estimate
+
+            else:
+                printer.log('Diamond distance wildcard model is incompatible with the Target estimate, skipping this.',3)
+                continue
 
         elif badfit_typ == "do nothing":
             continue  # go to next on-bad-fit directive
@@ -1869,6 +1914,182 @@ def _add_badfit_estimates(results, base_estimate_label, badfit_options,
                 results.estimates[base_estimate_label + '.' + badfit_typ].add_gaugeoptimized(
                     gauge_opt_params.copy(), go_gs_final, gokey, comm, printer - 1)
 
+def _compute_wildcard_budget_ddist_model(estimate, objfn_cache, mdc_objfn, parameters, badfit_options, verbosity):
+    """
+    Create a wildcard budget for a model estimate. This version of the function produces a wildcard estimate
+    using the model introduced by Tim and Stefan in the RCSGST paper.
+    TODO: docstring (update)
+
+    Parameters
+    ----------
+    model : Model
+        The model to add a wildcard budget to.
+
+    ds : DataSet
+        The data the model predictions are being compared with.
+
+    circuits_to_use : list
+        The circuits whose data are compared.
+
+    parameters : dict
+        Various parameters of the estimate at hand.
+
+    badfit_options : GSTBadFitOptions, optional
+        Options specifying what post-processing actions should be performed when
+        a fit is unsatisfactory.  Contains detailed parameters for wildcard budget
+        creation.
+
+    comm : mpi4py.MPI.Comm, optional
+        An MPI communicator used to run this computation in parallel.
+
+    mem_limit : int, optional
+        A rough per-processor memory limit in bytes.
+
+    verbosity : int, optional
+        Level of detail printed to stdout.
+
+    Returns
+    -------
+    PrimitiveOpsWildcardBudget
+    """
+    printer = _baseobjs.VerbosityPrinter.create_printer(verbosity, mdc_objfn.resource_alloc)
+    badfit_options = GSTBadFitOptions.cast(badfit_options)
+    model = mdc_objfn.model
+    ds = mdc_objfn.dataset
+    global_circuits_to_use = mdc_objfn.global_circuits
+
+    printer.log("******************* Adding Wildcard Budget **************************")
+
+    #cache construction code.
+    # Extract model, dataset, objfn, etc.
+    # Note: must evaluate mdc_objfn *before* passing to wildcard fn init so internal probs are init
+    mdc_objfn.fn(model.to_vector())
+
+    ## PYGSTI TRANSPLANT from pygsti.protocols.gst._compute_wildcard_budget
+    # Compute the various thresholds
+    ds_dof = ds.degrees_of_freedom(global_circuits_to_use)
+    nparams = model.num_modeltest_params  # just use total number of params
+    percentile = 0.025
+    nboxes = len(global_circuits_to_use)
+
+    two_dlogl_threshold = _chi2.ppf(1 - percentile, max(ds_dof - nparams, 1))
+    redbox_threshold  = _chi2.ppf(1 - percentile / nboxes, 1)
+
+    # Not in cache, recompute
+    ddists = _compute_ddists(estimate)
+
+    critical_percircuit_budgets = _opt.wildcardopt._get_critical_circuit_budgets(mdc_objfn, redbox_threshold)
+
+    #In Stefan's original code these were all passed in as a dictionary, instead pass these in as arguments.
+
+    alpha = _bisect_alpha(0.1, 1e-3, mdc_objfn, two_dlogl_threshold,
+                          redbox_threshold, ddists, critical_percircuit_budgets)
+
+    return _ddist_wildcard_model(alpha, ddists)
+
+def _bisect_alpha(guess, tol, mdc_objfn, two_dlogl_threshold, redbox_threshold, 
+                  ddists, critical_percircuit_budgets):
+
+    left = None
+    right = None
+
+    while left is None or right is None:
+        print(f'Searching for interval [{left}, {right}] with guess {guess}')
+        # Test for feasibility
+        guessval = _test_feasible_alphas([guess], mdc_objfn, two_dlogl_threshold, 
+                                        redbox_threshold, ddists, critical_percircuit_budgets)
+        if guessval[0]:
+            print('Guess value is feasible, ', end='')
+            left = guess
+            guess = left/2
+        else:
+            print('Guess value is infeasible, ', end='')
+            right = guess
+            guess = 2*right
+    print('Interval found!')
+
+    # We now have an interval containing the crossover point
+    # Perform bisection
+    while abs(left - right) > tol:
+        print(f'Performing bisection on interval [{left}, {right}]')
+        test = left - (left - right)/2.0
+
+        testval = _test_feasible_alphas([test], mdc_objfn, two_dlogl_threshold, 
+                                        redbox_threshold, ddists, critical_percircuit_budgets)
+
+        if testval[0]:
+            # Feasible, so shift left down
+            print(f'Test value is feasible, ', end='')
+            left = test
+        else:
+            print(f'Test value is infeasible, ', end='')
+            right = test
+
+    print('Interval within tolerance!')
+    print()
+
+    return left # Return the feasible one
+
+def _test_feasible_alphas(alphas, mdc_objfn, two_dlogl_threshold, redbox_threshold, ddists, critical_percircuit_budgets):
+
+    if not isinstance(alphas, list):
+        alphas = [alphas]
+
+    # Iterate through alphas and compute wildcard feasibilities
+    feasible = []
+    for alpha in alphas:
+        wcm = _ddist_wildcard_model(alpha, ddists)
+
+        ## PYGSTI TRANSPLANT from pygsti.optimize.wildcardopt.optimize_wildcard_budget_barrier
+        percircuit_budget_deriv = wcm.precompute_for_same_circuits(mdc_objfn.layout.circuits)
+
+        global_percircuit_budget_deriv_cols = []
+        for i in range(percircuit_budget_deriv.shape[1]):
+            global_percircuit_budget_deriv_cols.append(
+                mdc_objfn.layout.allgather_local_array('c', percircuit_budget_deriv[:, i]))
+        global_percircuit_budget_deriv = _np.column_stack(global_percircuit_budget_deriv_cols)
+
+        initial_probs = mdc_objfn.probs.copy()
+        current_probs = initial_probs.copy()
+        probs_freqs_precomp = wcm.precompute_for_same_probs_freqs(initial_probs, mdc_objfn.freqs, mdc_objfn.layout)
+
+        # Feasibility calculation, similar to penalty_vec function
+        x = wcm.to_vector()
+        wcm.update_probs(initial_probs, current_probs, mdc_objfn.freqs, mdc_objfn.layout, percircuit_budget_deriv,
+                            probs_freqs_precomp)
+        f0 = _np.array([_opt.wildcardopt._agg_dlogl(current_probs, mdc_objfn, two_dlogl_threshold)])
+        fi = critical_percircuit_budgets - _np.dot(global_percircuit_budget_deriv, x)
+
+        feasible.append(_np.all(_np.concatenate((f0, fi)) <= 0)) # All constraints must be negative to be feasible
+
+    return feasible
+
+def _compute_ddists(estimate):
+    final_model = estimate.models['final iteration estimate']
+    target_model = estimate.models['target']
+    gaugeopt_model = _alg.gaugeopt_to_target(final_model, target_model)
+
+    dd = {}
+    for key, op in gaugeopt_model.operations.items():
+        dd[key] = 0.5*_tools.diamonddist(op.to_dense(), target_model.operations[key].to_dense())
+
+    spamdd = {}
+    for key, op in gaugeopt_model.preps.items():
+        spamdd[key] = _tools.tracedist(_tools.vec_to_stdmx(op.to_dense(), 'pp'), _tools.vec_to_stdmx(target_model.preps[key].to_dense(), 'pp'))
+
+    for key in gaugeopt_model.povms.keys():
+        spamdd[key] = 0.5*_tools.optools.povm_diamonddist(gaugeopt_model, target_model, key)
+
+    dd['SPAM'] = sum(spamdd.values())
+
+    return dd
+
+def _ddist_wildcard_model(alpha, dd):
+    keys = list(dd.keys())
+    wcb_values = _np.array([alpha * dd[key] for key in keys])
+    wcb_dict= {key:value for key,value in zip(keys,wcb_values)}
+    wcm = _wild.PrimitiveOpsDiamondDistanceWildcardBudget(keys, start_budget=wcb_dict, ddists=dd, alpha=alpha) 
+    return wcm
 
 def _compute_robust_scaling(scale_typ, objfn_cache, mdc_objfn):
     """