Create a transformer that cancels the effect of Z-phases

cirq-core/cirq/experiments/z_phase_calibration.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """Provides a method to do z-phase calibration for excitation-preserving gates."""
-from typing import Union, Optional, Sequence, Tuple, Dict, TYPE_CHECKING, Any
+from typing import Union, Optional, Sequence, Tuple, Dict, TYPE_CHECKING, Any, List
 import multiprocessing
 import multiprocessing.pool
 
@@ -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, protocols
 
 if TYPE_CHECKING:
     import cirq
@@ -283,3 +284,84 @@ def plot_z_phase_calibration_result(
         ax.set_title('-'.join(str(q) for q in pair))
         ax.legend()
     return axes
+
+
+def _z_angles(old: ops.PhasedFSimGate, new: ops.PhasedFSimGate) -> Tuple[float, float, float]:
+    """Computes a set of possible 3 z-phases that result in the change in gamma, zeta, and chi."""
+    # This procedure is the inverse of PhasedFSimGate.from_fsim_rz
+    delta_gamma = new.gamma - old.gamma
+    delta_zeta = new.zeta - old.zeta
+    delta_chi = new.chi - old.chi
+    return (-delta_gamma + delta_chi, -delta_gamma - delta_zeta, delta_zeta - delta_chi)
+
+
+@transformer_api.transformer
+class CalibrationTransformer:
+
+    def __init__(
+        self,
+        target: 'cirq.Gate',
+        calibration_map: Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate'],
+    ):
+        """Create a CalibrationTransformer.
+
+        The transformer adds 3 ZPowGates around each calibrated gate to cancel the
+        effect of z-phases.
+
+        Args:
+            target: The target gate. Any gate matching this
+                will be replaced based on the content of `calibration_map`.
+            calibration_map:
+                A map mapping qubit pairs to calibrated gates. This is the output of
+                calling `calibrate_z_phases`.
+        """
+        self.target = target
+        if isinstance(target, ops.PhasedFSimGate):
+            self.target_as_fsim = target
+        elif (gate := ops.PhasedFSimGate.from_matrix(protocols.unitary(target))) is not None:
+            self.target_as_fsim = gate
+        else:
+            raise ValueError(f"{target} is not equivalent to a PhasedFSimGate")
+        self.calibration_map = calibration_map
+
+    def __call__(
+        self,
+        circuit: 'cirq.AbstractCircuit',
+        *,
+        context: Optional[transformer_api.TransformerContext] = None,
+    ) -> 'cirq.Circuit':
+        """Adds 3 ZPowGates around each calibrated gate to cancel the effect of Z phases.
+
+        Args:
+            circuit: Circuit to transform.
+            context: Optional transformer context (not used).
+
+        Returns:
+            New circuit with the extra ZPowGates.
+        """
+        new_moments: List[Union[List[cirq.Operation], 'cirq.Moment']] = []
+        for moment in circuit:
+            before = []
+            after = []
+            for op in moment:
+                if op.gate != self.target:
+                    # not a target.
+                    continue
+                assert len(op.qubits) == 2
+                gate = self.calibration_map.get(op.qubits, None) or self.calibration_map.get(
+                    op.qubits[::-1], None
+                )
+                if gate is None:
+                    # no calibration available.
+                    continue
+                angles = np.array(_z_angles(self.target_as_fsim, gate)) / np.pi
+                angles = -angles  # Take the negative to cancel the effect.
+                before.append(ops.Z(op.qubits[0]) ** angles[0])
+                before.append(ops.Z(op.qubits[1]) ** angles[1])
+                after.append(ops.Z(op.qubits[0]) ** angles[2])
+            if before:
+                new_moments.append(before)
+            new_moments.append(moment)
+            if after:
+                new_moments.append(after)
+        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,33 @@ 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])
+
+
+@pytest.mark.parametrize('angles', np.random.random((10, 10)))
+def test_transform_circuit(angles):
+    theta, phi = angles[:2]
+    old_zs = angles[2:6]
+    new_zs = angles[6:]
+    gate = cirq.PhasedFSimGate.from_fsim_rz(theta, phi, old_zs[:2], old_zs[2:])
+    fsim = cirq.PhasedFSimGate.from_fsim_rz(theta, phi, new_zs[:2], new_zs[2:])
+    c = cirq.Circuit(gate(cirq.q(0), cirq.q(1)))
+    replacement_map = {(cirq.q(1), cirq.q(0)): fsim}
+
+    new_circuit = CalibrationTransformer(gate, replacement_map)(c)
+
+    # we replace the old gate with the `fsim` gate the result should be that the overall
+    # unitary equals the unitary of the original (ideal) gate.
+    circuit_with_replacement_gate = cirq.Circuit(
+        op if op.gate != gate else fsim(*op.qubits) for op in new_circuit.all_operations()
+    )
+    np.testing.assert_allclose(cirq.unitary(circuit_with_replacement_gate), cirq.unitary(gate))
+
+
+def test_transform_circuit_invalid_gate_raises():
+    with pytest.raises(ValueError, match="is not equivalent to a PhasedFSimGate"):
+        _ = CalibrationTransformer(cirq.XX, {})
+
+
+def test_transform_circuit_uncalibrated_gates_pass():
+    c = cirq.Circuit(cirq.CZ(cirq.q(0), cirq.q(1)), cirq.measure(cirq.q(0)))
+    assert c == CalibrationTransformer(cirq.CZ, {})(c)

cirq-core/cirq/ops/fsim_gate.py

@@ -347,6 +347,35 @@ def from_fsim_rz(
         chi = (b0 - b1 - a0 + a1) / 2.0
         return PhasedFSimGate(theta, zeta, chi, gamma, phi)
 
+    @staticmethod
+    def from_matrix(u: np.ndarray) -> Optional['PhasedFSimGate']:
+        gamma = np.angle(u[1, 1] * u[2, 2] - u[1, 2] * u[2, 1]) / -2
+        phi = -np.angle(u[3, 3]) - 2 * gamma
+        phased_cos_theta_2 = u[1, 1] * u[2, 2]
+        if phased_cos_theta_2 == 0:
+            # The zeta phase is multiplied with cos(theta),
+            # so if cos(theta) is zero then any value is possible.
+            zeta = 0
+        else:
+            zeta = np.angle(u[2, 2] / u[1, 1]) / 2
+
+        phased_sin_theta_2 = u[1, 2] * u[2, 1]
+        if phased_sin_theta_2 == 0:
+            # The chi phase is multiplied with sin(theta),
+            # so if sin(theta) is zero then any value is possible.
+            chi = 0
+        else:
+            chi = np.angle(u[1, 2] / u[2, 1]) / 2
+
+        theta = np.angle(
+            np.exp(1j * (gamma + zeta)) * u[1, 1] - np.exp(1j * (gamma - chi)) * u[1, 2]
+        )
+
+        gate = PhasedFSimGate(theta=theta, phi=phi, chi=chi, zeta=zeta, gamma=gamma)
+        if np.allclose(u, protocols.unitary(gate)):
+            return gate
+        return None
+
     @property
     def rz_angles_before(self) -> Tuple['cirq.TParamVal', 'cirq.TParamVal']:
         """Returns 2-tuple of phase angles applied to qubits before FSimGate."""

cirq-core/cirq/ops/fsim_gate_test.py

@@ -797,3 +797,24 @@ def test_phased_fsim_json_dict():
     assert cirq.PhasedFSimGate(
         theta=0.12, zeta=0.34, chi=0.56, gamma=0.78, phi=0.9
     )._json_dict_() == {'theta': 0.12, 'zeta': 0.34, 'chi': 0.56, 'gamma': 0.78, 'phi': 0.9}
+
+
+@pytest.mark.parametrize(
+    'gate',
+    [
+        cirq.CZ,
+        cirq.SQRT_ISWAP,
+        cirq.SQRT_ISWAP_INV,
+        cirq.ISWAP,
+        cirq.ISWAP_INV,
+        cirq.cphase(0.1),
+        cirq.CZ**0.2,
+    ],
+)
+def test_phase_fsim_from_matrix(gate):
+    u = cirq.unitary(gate)
+    np.testing.assert_allclose(cirq.unitary(cirq.PhasedFSimGate.from_matrix(u)), u, atol=1e-8)
+
+
+def test_phase_fsim_from_matrix_not_fsim_returns_none():
+    assert cirq.PhasedFSimGate.from_matrix(np.ones((4, 4))) is None