Parameterization-preserving gauge optimization

[rileyjmurray]

EDIT: versions of the functions below were merged in as part of PR #410. I'm sure work remains on this issue (e.g., points raised by Corey below) so I'm leaving it open.

---

Consider the following fact.

Let $M_1$ be a model (e.g., an ExplicitOpModel) with parameterization $P_1$ (e.g., "CPTPLND" or "full TP"). Next, let $M_2$ be the representation of $M_1$ under a parameterization $P_2$, where $P_2$ allows for a strictly larger hypothesis class than does $P_1$.

We can't, in general, get confidence intervals for $M_1$ by looking at the likelihood function for $M_2$.

This fact is troublesome for computing confidence intervals of gauge-dependent quantities, since pyGSTi's gauge optimization tends to be performed in an expanded hypothesis class. This is the case even if we do gauge optimization on a CPTP model (CPTPLND parameterization) over the unitary group.

It's nominally straightforward to avoid this issue by gauge optimizing in whatever parameterization we want (e.g., Full), getting the gauge transformation, and then substituting modelmembers with suitably-transformed versions of themselves. The transformation + substitution can be facilitated with ComposedState, ComposedOp, and ComposedPOVM, as shown in the following code.

def transform_composedop_model(mdl : ExplicitOpModel, s : GaugeGroupElement):
    oldmdl = mdl
    mdl = oldmdl.copy()

    U    = StaticArbitraryOp(s.transform_matrix,         basis=oldmdl.basis)
    invU = StaticArbitraryOp(s.transform_matrix_inverse, basis=oldmdl.basis) 

    for key, rho in oldmdl.preps.items():
        assert isinstance(rho, ComposedState)
        static_rho = rho.state_vec
        errmap  = ComposedOp([rho.error_map, invU])
        mdl.preps[key] = ComposedState(static_rho, errmap)

    for key, povm in oldmdl.povms.items():
        assert isinstance(povm, ComposedPOVM)
        static_povm = povm.base_povm
        errmap = ComposedOp([U, povm.error_map])
        mdl.povms[key] = ComposedPOVM(errmap, static_povm, mx_basis=oldmdl.basis)

    for key, op in oldmdl.operations.items():
        op_s = ComposedOp([U, op, invU])
        mdl.operations[key] = op_s

    assert len(oldmdl.factories) == 0
    assert len(oldmdl.instruments) == 0

    mdl._clean_paramvec()  # transform may leave dirty members
    return mdl

A whole gauge optimization pipeline could look something like this.

def add_param_preserving_gauge_opt(results: ModelEstimateResults, est_key: str, gop_params: GSTGaugeOptSuite):
    from pygsti.protocols.gst import _add_gauge_opt
    est = results.estimates[est_key]
    seed_mdl = est.models['final iteration estimate']
    _add_gauge_opt(results, est_key, gop_params, seed_mdl)
    # ^ That can convert to whatever parameterization it wants. It'll write to 
    #   est._gaugeopt_suite.
    gop_params_dict = gop_params.to_dictionary(seed_mdl)  
    # ^ If we modified _add_gauge_opt to accept both GSTGaugeOptSuite objects and
    #   the output of the .to_dictionary(...) method of such objects, then we could avoid
    #   this awkward call to .to_dictionary(...) after the call to _add_gauge_opt.
    for gop_name in gop_params_dict.keys():
        ggel = est._gaugeopt_suite.gaugeopt_argument_dicts[gop_name]['_gaugeGroupEl']
        model_implicit_gauge = transform_composed_model(est.models['final iteration estimate'], ggel)
        est.models[gop_name] = model_implicit_gauge
    return

I want to integrate this into pyGSTi's default gauge optimization procedures that get executed when run(...) is called on a GateSetTomography object. For example, maybe we allow passing gaugeopt_suite='stdgaugeopt-keep-params' to the GateSetTomography constructor, and we get behavior that's just like if we used gaugeopt_suite='stdgaugeopt' without having to accept the change in model parameterization.

Here are some observations for myself as I try to sort this out.

• GateSetTomography.run calls _add_gaugeopt_and_badfit, which calls _add_gauge_opt and _add_badfit_estimates. Meanwhile, _add_badfit_estimates can call _add_gauge_opt.
• If we look at _add_gauge_opt we see that it takes in a GSTGaugeOptSuite object, which it immediately converts to a dict representation with a call to the object's .to_dictionary function (note: that function requires that we specify a starting model for gauge optimization, since no models are stored in GSTGaugeOptSuite objects). Then _add_gauge_opt makes calls to the add_gaugeoptimized instance function of Estimate objects from the .estimates field of results (which is a ModelEstimateResults object).

[coreyostrove]

Relevant thing that I figured out/ran into this afternoon. I recently found/fixed some out-of-date logic for setting the default gauge groups for models when using set_all_parameterizations. This is the commit in question 8307cd5. TLDR: before lindblad parameterizations were getting their default gauge group set to the trivial gauge group, and I fixed/updated this to be the unitary group as one would more likely expect. In doing so I inadvertently re-activated a latent code path for applying gauge transformations to ComposedOp and ComposedState etc, which revealed there are some bugs and/or unexpected behavior which were flagged by some unit tests that started failing. I couldn't immediately pin this down, and in the meantime have implemented a workaround, but in the course of trying to implement your idea here we'll probably need to figure this out. Below is some minimal code (adapted from one of the unit tests which caught this) for reproducing the error, along with the corresponding traceback.

import pygsti
from pygsti.modelpacks import smq1Q_XY as std
model = std.target_model()
pspec = std.processor_spec()
opLabels = list(model.operations.keys())
prep_fids = std.prep_fiducials()
meas_fids = std.meas_fiducials()
germs = std.germs(lite=True)
maxLengthList = [1, 2, 4]
datagen_model = model.depolarize(op_noise=0.05, spam_noise=0.1)
circuits = pygsti.circuits.create_lsgst_circuit_lists(
        opLabels, prep_fids, meas_fids,
        germs, maxLengthList
    )
ds = pygsti.data.simulate_data(
        datagen_model, circuits[-1],
        num_samples=1000, sample_error='binomial', seed=100
    )
model.set_all_parameterizations("CPTPLND")
result = pygsti.drivers.longsequence.run_long_sequence_gst(
            ds, model, prep_fids, meas_fids,
            germs, maxLengthList)

Which results in this traceback:

 ---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[3], line 2
      1 model.set_all_parameterizations("CPTPLND")
----> 2 result = pygsti.drivers.longsequence.run_long_sequence_gst(
      3             ds, model, prep_fids, meas_fids,
      4             germs, maxLengthList)

File ~\Documents\pyGSTi_random_bugfixes\pygsti\drivers\longsequence.py:509, in run_long_sequence_gst(data_filename_or_set, target_model_filename_or_object, prep_fiducial_list_or_filename, meas_fiducial_list_or_filename, germs_list_or_filename, max_lengths, gauge_opt_params, advanced_options, comm, mem_limit, output_pkl, verbosity, checkpoint, checkpoint_path, disable_checkpointing, simulator)
    506 proto.circuit_weights = advanced_options.get('circuit_weights', None)
    507 proto.unreliable_ops = advanced_options.get('unreliable_ops', ['Gcnot', 'Gcphase', 'Gms', 'Gcn', 'Gcx', 'Gcz'])
--> 509 results = proto.run(data, mem_limit, comm,
    510                     checkpoint=checkpoint, checkpoint_path= checkpoint_path, disable_checkpointing=disable_checkpointing,
    511                     simulator=simulator)
    512 _output_to_pickle(results, output_pkl, comm)
    513 return results

File ~\Documents\pyGSTi_random_bugfixes\pygsti\protocols\gst.py:1495, in GateSetTomography.run(self, data, memlimit, comm, checkpoint, checkpoint_path, disable_checkpointing, simulator)
   1493 #Add some better handling for when gauge optimization is turned off (current code path isn't working.)
   1494 if not self.gaugeopt_suite.is_empty() or len(self.badfit_options.actions) > 0:  # maybe add flag to do this even when empty?
-> 1495     ret = _add_gaugeopt_and_badfit(ret, self.name, target_model,
   1496                                    self.gaugeopt_suite, self.unreliable_ops,
   1497                                    self.badfit_options, self.optimizer, 
   1498                                    resource_alloc, printer)
   1499 else:
   1500     #add a model to the estimate that we'll call the trivial gauge optimized model which
   1501     #will be set to be equal to the final iteration estimate.
   1502     ret.estimates[self.name].models['trivial_gauge_opt'] = mdl_lsgst_list[-1]

File ~\Documents\pyGSTi_random_bugfixes\pygsti\protocols\gst.py:2041, in _add_gaugeopt_and_badfit(results, estlbl, target_model, gaugeopt_suite, unreliable_ops, badfit_options, optimizer, resource_alloc, printer)
   2039     gaugeopt_suite = _copy.deepcopy(gaugeopt_suite)
   2040     gaugeopt_suite.gaugeopt_target = target_model
-> 2041 _add_gauge_opt(results, estlbl, gaugeopt_suite,
   2042                model_to_gaugeopt, unreliable_ops, comm, printer - 1)
   2043 profiler.add_time('%s: gauge optimization' % estlbl, tref); tref = _time.time()
   2045 _add_badfit_estimates(results, estlbl, badfit_options, optimizer, resource_alloc, printer, gaugeopt_suite= gaugeopt_suite)

File ~\Documents\pyGSTi_random_bugfixes\pygsti\protocols\gst.py:2117, in _add_gauge_opt(results, base_est_label, gaugeopt_suite, starting_model, unreliable_ops, comm, verbosity)
   2115     results.estimates[base_est_label].add_gaugeoptimized(goparams, starting_model, go_label, comm, printer - 3)
   2116 else:
-> 2117     results.estimates[base_est_label].add_gaugeoptimized(goparams, None, go_label, comm, printer - 3)
   2119 #Get starting model for next stage
   2120 mdl_start = results.estimates[base_est_label].retrieve_start_model(goparams)

File ~\Documents\pyGSTi_random_bugfixes\pygsti\protocols\estimate.py:405, in Estimate.add_gaugeoptimized(self, goparams, model, label, comm, verbosity)
    402 gop['return_all'] = True
    403 if isinstance(gop['model'], _ExplicitOpModel):
    404     #only explicit models can be gauge optimized
--> 405     _, gauge_group_el, last_gs = _gaugeopt_to_target(**gop)
    406 else:
    407     #but still fill in results for other models (?)
    408     gauge_group_el, last_gs = None, gop['model'].copy()

File ~\Documents\pyGSTi_random_bugfixes\pygsti\algorithms\gaugeopt.py:175, in gaugeopt_to_target(model, target_model, item_weights, cptp_penalty_factor, spam_penalty_factor, gates_metric, spam_metric, gauge_group, method, maxiter, maxfev, tol, oob_check_interval, convert_model_to, return_all, comm, verbosity, check_jac)
    168             raise ValueError("Invalid `convert_model_to` arguments: %s" % str(args))
    170 objective_fn, jacobian_fn = _create_objective_fn(
    171     model, target_model, item_weights,
    172     cptp_penalty_factor, spam_penalty_factor,
    173     gates_metric, spam_metric, method, comm, check_jac)
--> 175 result = gaugeopt_custom(model, objective_fn, gauge_group, method,
    176                          maxiter, maxfev, tol, oob_check_interval,
    177                          return_all, jacobian_fn, comm, verbosity)
    179 #If we've gauge optimized to a target model, declare that the
    180 # resulting model is now in the same basis as the target.
    181 if target_model is not None:

File ~\Documents\pyGSTi_random_bugfixes\pygsti\algorithms\gaugeopt.py:316, in gaugeopt_custom(model, objective_fn, gauge_group, method, maxiter, maxfev, tol, oob_check_interval, return_all, jacobian_fn, comm, verbosity)
    314 ralloc = _baseobjs.ResourceAllocation(comm)  # FUTURE: plumb up a resource alloc object?
    315 test_f = _call_objective_fn(x0)
--> 316 solnX, converged, msg, _, _, _, _ = _opt.simplish_leastsq(
    317     _call_objective_fn, _call_jacobian_fn, x0, f_norm2_tol=tol,
    318     jac_norm_tol=tol, rel_ftol=tol, rel_xtol=tol,
    319     max_iter=maxiter, resource_alloc=ralloc,
    320     arrays_interface=_opt.UndistributedArraysInterface(len(test_f), len(x0)),
    321     oob_check_interval=oob_check_interval,
    322     verbosity=printer.verbosity - 2)
    323 printer.log("Least squares message = %s" % msg, 2)
    324 assert(converged)

File ~\Documents\pyGSTi_random_bugfixes\pygsti\optimize\simplerlm.py:709, in simplish_leastsq(obj_fn, jac_fn, x0, f_norm2_tol, jac_norm_tol, rel_ftol, rel_xtol, max_iter, num_fd_iters, max_dx_scale, init_munu, oob_check_interval, oob_action, oob_check_mode, resource_alloc, arrays_interface, serial_solve_proc_threshold, x_limits, verbosity, profiler)
    706 else:
    707     #Just evaluate objective function normally; never check for in-bounds condition
    708     ari.allgather_x(new_x, global_new_x)
--> 709     new_f = obj_fn(global_new_x)
    710     new_x_is_known_inbounds = oob_check_interval == 0
    711     # ^ assume in bounds if we have no out-of-bounds checks.

File ~\Documents\pyGSTi_random_bugfixes\pygsti\algorithms\gaugeopt.py:299, in gaugeopt_custom.<locals>._call_objective_fn(gauge_group_el_vec, oob_check)
    296 def _call_objective_fn(gauge_group_el_vec, oob_check=False):
    297     # Note: oob_check can be True if oob_check_interval>=1 is given to the simplish_leastsq below
    298     gaugeGroupEl.from_vector(gauge_group_el_vec)
--> 299     return objective_fn(gaugeGroupEl, oob_check)

File ~\Documents\pyGSTi_random_bugfixes\pygsti\algorithms\gaugeopt.py:408, in _create_objective_fn.<locals>._objective_fn(gauge_group_el, oob_check)
    406     other = model
    407 else:
--> 408     transformed = _transform_with_oob_check(model, gauge_group_el, oob_check)
    409     other = target_model
    411 residuals, _ = transformed.residuals(other, None, item_weights)

File ~\Documents\pyGSTi_random_bugfixes\pygsti\algorithms\gaugeopt.py:384, in _create_objective_fn.<locals>._transform_with_oob_check(mdl, gauge_group_el, oob_check)
    382         raise ValueError("Out of bounds: %s" % str(e))  # signals OOB condition
    383 else:
--> 384     mdl.transform_inplace(gauge_group_el)
    385 return mdl

File ~\Documents\pyGSTi_random_bugfixes\pygsti\models\explicitmodel.py:688, in ExplicitOpModel.transform_inplace(self, s)
    671 """
    672 Gauge transform this model.
    673 
   (...)
    685 None
    686 """
    687 for rhoVec in self.preps.values():
--> 688     rhoVec.transform_inplace(s)
    690 for povm in self.povms.values():
    691     povm.transform_inplace(s)

File ~\Documents\pyGSTi_random_bugfixes\pygsti\modelmembers\states\composedstate.py:420, in ComposedState.transform_inplace(self, s)
    394 """
    395 Update state (column) vector V as inv(s) * V or s^T * V for preparation or  effect state vectors, respectively.
    396 
   (...)
    414 None
    415 """
    416 #Defer everything to LindbladOp's
    417 # `spam_tranform` function, which applies either
    418 # error_map -> inv(s) * error_map ("prep" case) OR
    419 # error_map -> error_map * s      ("effect" case)
--> 420 self.error_map.spam_transform_inplace(s, 'prep')
    421 self._update_rep()
    422 self.dirty = True

File ~\Documents\pyGSTi_random_bugfixes\pygsti\modelmembers\operations\experrorgenop.py:706, in ExpErrorgenOp.spam_transform_inplace(self, s, typ)
    704     else:
    705         mx = mx @ U
--> 706     self.set_dense(mx)  # calls _update_rep() and sets dirty flag
    707 else:
    708     raise ValueError("Invalid transform for this LindbladErrorgen: type %s"
    709                      % str(type(s)))

File ~\Documents\pyGSTi_random_bugfixes\pygsti\modelmembers\operations\experrorgenop.py:635, in ExpErrorgenOp.set_dense(self, m)
    633 #Note: this only really works for LindbladErrorGen objects now... make more general in FUTURE?
    634 truncate = TODENSE_TRUNCATE  # can't just be 'True' since we need to throw errors when appropriate
--> 635 new_errgen = errgen_cls.from_operation_matrix_and_blocks(
    636     mx, self.errorgen.coefficient_blocks, 'auto', self.errorgen.matrix_basis,
    637     truncate, self.errorgen.evotype, self.errorgen.state_space)
    638 self.errorgen.from_vector(new_errgen.to_vector())
    639 self._update_rep()  # needed to rebuild exponentiated error gen

File ~\Documents\pyGSTi_random_bugfixes\pygsti\modelmembers\operations\lindbladerrorgen.py:119, in LindbladErrorgen.from_operation_matrix_and_blocks(cls, op_matrix, lindblad_coefficient_blocks, elementary_errorgen_basis, mx_basis, truncate, evotype, state_space)
    116     errgenMx = _ot.error_generator(op_matrix, _np.identity(op_matrix.shape[0], 'd'),
    117                                    mx_basis, "logGTi")
    118 for blk in lindblad_coefficient_blocks:
--> 119     blk.set_from_errorgen_projections(errgenMx, mx_basis, truncate=truncate)
    120 return cls(lindblad_coefficient_blocks, elementary_errorgen_basis, mx_basis, evotype, state_space)

File ~\Documents\pyGSTi_random_bugfixes\pygsti\modelmembers\operations\lindbladcoefficients.py:547, in LindbladCoefficientBlock.set_from_errorgen_projections(self, errorgen, errorgen_basis, return_projected_errorgen, truncate)
    543 out = _ot.extract_elementary_errorgen_coefficients(errorgen, elementary_errorgen_lbls,
    544                                                    self._basis, errorgen_basis,
    545                                                    return_projected_errorgen)
    546 elementary_errorgens = out[0] if return_projected_errorgen else out
--> 547 unused = self.set_elementary_errorgens(elementary_errorgens, on_missing='raise', truncate=truncate)
    548 assert(len(unused) == 0)
    550 #set a flag to indicate that the coefficients (as returned by elementary_errorgens)
    551 #need to be updated.

File ~\Documents\pyGSTi_random_bugfixes\pygsti\modelmembers\operations\lindbladcoefficients.py:532, in LindbladCoefficientBlock.set_elementary_errorgens(self, elementary_errorgens, on_missing, truncate)
    529     flat_data[i] = val
    531 self.block_data[:] = flat_data.reshape(self.block_data.shape)
--> 532 self._truncate_block_data(truncate)
    534 #set a flag to indicate that the coefficients (as returned by elementary_errorgens)
    535 #need to be updated.
    536 self._coefficients_need_update = True

File ~\Documents\pyGSTi_random_bugfixes\pygsti\modelmembers\operations\lindbladcoefficients.py:782, in LindbladCoefficientBlock._truncate_block_data(self, truncate)
    780 def _truncate_block_data(self, truncate=False):
    781     if truncate is False: return
--> 782     v = self._block_data_to_params(truncate)
    783     self.from_vector(v)

File ~\Documents\pyGSTi_random_bugfixes\pygsti\modelmembers\operations\lindbladcoefficients.py:730, in LindbladCoefficientBlock._block_data_to_params(self, truncate)
    726 assert(_np.isclose(_np.linalg.norm(self.block_data), _np.linalg.norm(nonzero_block_data)))
    728 evals, U = _np.linalg.eigh(nonzero_block_data)
--> 730 assert(all([ev > ttol for ev in evals])), \
    731     ("Lindblad coefficients are not CPTP (truncate == %s)! (largest neg = %g)"
    732      % (str(truncate), min(evals)))
    734 if ttol < 0:  # if we're truncating and assert above allows *negative* eigenvalues
    735     #push any slightly negative evals of other_projs positive so that
    736     # the Cholesky decomp will work.
    737     Ui = U.T.conj()

AssertionError: Lindblad coefficients are not CPTP (truncate == 3e-10)! (largest neg = -1.11683e-06)

I did trace through the logic for doing the gauge transformation on the state preparation's corresponding error generator (as it has when using the CPTPLND parameterization), and while I didn't spot any obvious mistakes in the logic, there must be something subtle going on. The logic raising the error performs a gauge transformation on the SPAM by multiplying the corresponding error channel (i.e. exp(L_rho)) by a unitary matrix and then attempting to construct a new error generator by taking the logarithm and doing projections. This resulting error generator is getting flagged as non-CP, and the block_data for the SCA block is showing up with negative eigenvalues as the traceback says, but the corresponding error channel itself is CP so far as I can tell (it has no negative choi eigenvalues).

I'm now thinking out loud, but it is plausible that we're running into one of those distinctions between CP and infinitesimally-CP-divisible. I.e. the manner in which we're applying the gauge transformation may indeed by resulting in a CP channel which nonetheless cannot be represented by a CP error generator.

The use of the function run_longsequence_gst for this is important in this case, as it had default settings configured which, unlike our more commonly used GST code paths, didn't perform a conversion to full/full TP before gauge optimization (thus letting it hit the code path above). I have since updated the default behavior of run_long_sequence_gst to use the 'stdgaugeopt' suite, which does include this conversion (this was an independently desirable change, and brings this function in line with its OOP protocol analogue of GateSetTomography). So to reproduce this make sure you're on the aforementioned commit.