From 5eebdc0e72d2c925bd6d35893afd3afcdefabae0 Mon Sep 17 00:00:00 2001 From: Corey Ostrove Date: Fri, 20 Jun 2025 17:51:32 -0600 Subject: [PATCH] Add in missing return_basis implementation Patch in support for the missing support for the return_basis of EmbeddedErrorgen and EmbeddedOp. --- .../modelmembers/operations/embeddederrorgen.py | 16 +++++++++++++--- pygsti/modelmembers/operations/embeddedop.py | 16 +++++++++++++--- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/pygsti/modelmembers/operations/embeddederrorgen.py b/pygsti/modelmembers/operations/embeddederrorgen.py index 47614899e..a8a3c6a1b 100644 --- a/pygsti/modelmembers/operations/embeddederrorgen.py +++ b/pygsti/modelmembers/operations/embeddederrorgen.py @@ -139,16 +139,26 @@ def coefficients(self, return_basis=False, logscale_nonham=False, label_type='gl Where `termType` is `"H"` (Hamiltonian), `"S"` (Stochastic), `"C"`(Correlation) or `"A"` (Affine). Hamiltonian and S terms always have a single basis label while 'C' and 'A' terms have two. - """ - coeffs_to_embed = self.embedded_op.coefficients(return_basis, logscale_nonham, label_type) + basis : `Basis` (if return_basis==True) + A Basis mapping the basis labels used in the keys of basis-labels of the + underlying (pre-embedding) error generator coeffs to basis matrices. + """ + if return_basis: + coeffs_to_embed, basis = self.embedded_op.coefficients(return_basis, logscale_nonham, label_type) + else: + coeffs_to_embed = self.embedded_op.coefficients(return_basis, logscale_nonham, label_type) + if coeffs_to_embed: embedded_labels = self.coefficient_labels(label_type=label_type, identity_label=identity_label) embedded_coeffs = {lbl:val for lbl, val in zip(embedded_labels, coeffs_to_embed.values())} else: embedded_coeffs = dict() - return embedded_coeffs + if return_basis: + return embedded_coeffs, basis + else: + return embedded_coeffs def coefficient_labels(self, label_type='global', identity_label='I'): """ diff --git a/pygsti/modelmembers/operations/embeddedop.py b/pygsti/modelmembers/operations/embeddedop.py index 980038fbb..2b7f1f942 100644 --- a/pygsti/modelmembers/operations/embeddedop.py +++ b/pygsti/modelmembers/operations/embeddedop.py @@ -600,17 +600,27 @@ def errorgen_coefficients(self, return_basis=False, logscale_nonham=False, label Where `termType` is `"H"` (Hamiltonian), `"S"` (Stochastic), `"C"`(Correlation) or `"A"` (Affine). Hamiltonian and S terms always have a single basis label while 'C' and 'A' terms have two. + + basis : `Basis` (if return_basis==True) + A Basis mapping the basis labels used in the keys of basis-labels of the + underlying (pre-embedding) error generator coeffs to basis matrices. """ #*** Note: this function is nearly identical to EmbeddedErrorgen.coefficients() *** - coeffs_to_embed = self.embedded_op.errorgen_coefficients(return_basis, logscale_nonham, label_type) - + if return_basis: + coeffs_to_embed, basis = self.embedded_op.errorgen_coefficients(return_basis, logscale_nonham, label_type) + else: + coeffs_to_embed, basis = self.embedded_op.errorgen_coefficients(return_basis, logscale_nonham, label_type) + if coeffs_to_embed: embedded_labels = self.errorgen_coefficient_labels(label_type=label_type, identity_label=identity_label) embedded_coeffs = {lbl:val for lbl, val in zip(embedded_labels, coeffs_to_embed.values())} else: embedded_coeffs = dict() - return embedded_coeffs + if return_basis: + return embedded_coeffs, basis + else: + return embedded_coeffs def errorgen_coefficient_labels(self, label_type='global', identity_label='I'): """