Create a transformer cancels the effect of Z-phases

cirq-core/cirq/experiments/z_phase_calibration.py

@@ -22,7 +22,8 @@
 
 from cirq.experiments import xeb_fitting
 from cirq.experiments.two_qubit_xeb import parallel_xeb_workflow
-from cirq import ops
+from cirq.transformers import transformer_api
+from cirq import ops, circuits
 
 if TYPE_CHECKING:
     import cirq
@@ -271,3 +272,59 @@ def plot_z_phase_calibration_result(
         ax.set_title('-'.join(str(q) for q in pair))
         ax.legend()
     return axes
+
+
+@transformer_api.transformer
+class CalibrationTransformer:
+
+    def __init__(
+        self,
+        target: Union['cirq.Gate', 'cirq.GateFamily', 'cirq.Gateset'],
+        replacement_map: Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate'],
+    ):
+        """Create a CalibrationTransformer that replaces gates matching `target`.
+
+        Args:
+            target: The target gate, GateFamily or Gateset. Any gate matching this
+                will be replaced based on the content of `replacement_map`.
+            replacement_map:
+                A map mapping qubit pairs to calibrated gates. This is the output of
+                calling `calibrate_z_phases`.
+        """
+        self.target = ops.GateFamily(target) if isinstance(target, ops.Gate) else target
+        self.replacement_map = replacement_map
+
+    def __call__(
+        self,
+        circuit: 'cirq.AbstractCircuit',
+        *,
+        context: Optional[transformer_api.TransformerContext] = None,
+    ) -> 'cirq.Circuit':
+        """Replace every occurance of a calibrated gate with a proper replacement.
+
+        Args:
+            circuit: Circuit to transform.
+            context: Optional transformer context (not used).
+
+        Returns:
+            New circuit with each gate matching `target` and acting on a pair
+                in `replacement_map` replaced by the gate pointed to by `replacement_map`.
+        """
+        new_moments = []
+        for moment in circuit:
+            new_moment = []
+            for op in moment:
+                if op not in self.target:
+                    # not a target.
+                    new_moment.append(op)
+                    continue
+                gate = self.replacement_map.get(op.qubits, None) or self.replacement_map.get(
+                    op.qubits[::-1], None
+                )
+                if gate is None:
+                    # no calibrated version exist
+                    new_moment.append(op)
+                    continue
+                new_moment.append(gate(*op.qubits))
+            new_moments.append(new_moment)
+        return circuits.Circuit.from_moments(*new_moments)

cirq-core/cirq/experiments/z_phase_calibration_test.py

@@ -22,6 +22,7 @@
     calibrate_z_phases,
     z_phase_calibration_workflow,
     plot_z_phase_calibration_result,
+    CalibrationTransformer,
 )
 from cirq.experiments.xeb_fitting import XEBPhasedFSimCharacterizationOptions
 
@@ -205,3 +206,10 @@ def test_plot_z_phase_calibration_result():
     np.testing.assert_allclose(axes[1].lines[0].get_xdata().astype(float), [1, 2, 3])
     np.testing.assert_allclose(axes[1].lines[0].get_ydata().astype(float), [0.6, 0.4, 0.1])
     np.testing.assert_allclose(axes[1].lines[1].get_ydata().astype(float), [0.7, 0.77, 0.8])
+
+
+def test_transform_circuit():
+    c = cirq.Circuit(cirq.CZ(cirq.q(0), cirq.q(1)))
+    replacement_map = {(cirq.q(1), cirq.q(0)): cirq.PhasedFSimGate(0, 1, 2, 3, 4)}
+    new_circuit = CalibrationTransformer(cirq.CZ, replacement_map)(c)
+    assert new_circuit == cirq.Circuit(cirq.PhasedFSimGate(0, 1, 2, 3, 4).on(cirq.q(0), cirq.q(1)))