From ff7297f38e3b226faf69c3b9e0ae3133b7c543bd Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 11 Dec 2024 14:25:54 -0800
Subject: [PATCH 01/10] Create a transformer to replace gates with their
 calibrated versions

---
 .../cirq/experiments/z_phase_calibration.py   | 49 ++++++++++++++++++-
 .../experiments/z_phase_calibration_test.py   |  8 +++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 363810fbd13..28740a7ce82 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/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,49 @@ def plot_z_phase_calibration_result(
         ax.set_title('-'.join(str(q) for q in pair))
         ax.legend()
     return axes
+
+
+@transformer_api.transformer
+def transform_circuit(
+    circuit: 'cirq.AbstractCircuit',
+    target: Union['cirq.Gate', 'cirq.GateFamily', 'cirq.Gateset'],
+    replacement_map: Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate'],
+    *,
+    context: Optional[transformer_api.TransformerContext] = None,
+) -> 'cirq.Circuit':
+    """Replace every occurance of a calibrated gate with a proper replacement.
+
+    Args:
+        circuit: Circuit to transform.
+        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`.
+        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`.
+    """
+    if isinstance(target, ops.Gate):
+        target = ops.GateFamily(target)
+
+    new_moments = []
+    for moment in circuit:
+        new_moment = []
+        for op in moment:
+            if op not in target:
+                # not a target.
+                new_moment.append(op)
+                continue
+            gate = replacement_map.get(op.qubits, None) or 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)
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index c7c149b37c5..730c45d2c29 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/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,
+    transform_circuit,
 )
 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 = transform_circuit(circuit=c, target=cirq.CZ, replacement_map=replacement_map)
+    assert new_circuit == cirq.Circuit(cirq.PhasedFSimGate(0, 1, 2, 3, 4).on(cirq.q(0), cirq.q(1)))

From 9fe78a9b3762c1e9ae16d2c034e48f20ec0b4629 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 11 Dec 2024 14:32:55 -0800
Subject: [PATCH 02/10] types

---
 cirq-core/cirq/experiments/z_phase_calibration.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 28740a7ce82..fbd5b2ec641 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -277,9 +277,9 @@ def plot_z_phase_calibration_result(
 @transformer_api.transformer
 def transform_circuit(
     circuit: 'cirq.AbstractCircuit',
+    *,
     target: Union['cirq.Gate', 'cirq.GateFamily', 'cirq.Gateset'],
     replacement_map: Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate'],
-    *,
     context: Optional[transformer_api.TransformerContext] = None,
 ) -> 'cirq.Circuit':
     """Replace every occurance of a calibrated gate with a proper replacement.

From 801e14e8775c89b1de5e7304da1a2b70b38c2279 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 11 Dec 2024 14:41:25 -0800
Subject: [PATCH 03/10] make class

---
 .../cirq/experiments/z_phase_calibration.py   | 96 ++++++++++---------
 .../experiments/z_phase_calibration_test.py   |  4 +-
 2 files changed, 55 insertions(+), 45 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index fbd5b2ec641..d67f552932e 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -275,46 +275,56 @@ def plot_z_phase_calibration_result(
 
 
 @transformer_api.transformer
-def transform_circuit(
-    circuit: 'cirq.AbstractCircuit',
-    *,
-    target: Union['cirq.Gate', 'cirq.GateFamily', 'cirq.Gateset'],
-    replacement_map: Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate'],
-    context: Optional[transformer_api.TransformerContext] = None,
-) -> 'cirq.Circuit':
-    """Replace every occurance of a calibrated gate with a proper replacement.
-
-    Args:
-        circuit: Circuit to transform.
-        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`.
-        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`.
-    """
-    if isinstance(target, ops.Gate):
-        target = ops.GateFamily(target)
-
-    new_moments = []
-    for moment in circuit:
-        new_moment = []
-        for op in moment:
-            if op not in target:
-                # not a target.
-                new_moment.append(op)
-                continue
-            gate = replacement_map.get(op.qubits, None) or 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)
+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)
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 730c45d2c29..8ad90c83b95 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -22,7 +22,7 @@
     calibrate_z_phases,
     z_phase_calibration_workflow,
     plot_z_phase_calibration_result,
-    transform_circuit,
+    CalibrationTransformer,
 )
 from cirq.experiments.xeb_fitting import XEBPhasedFSimCharacterizationOptions
 
@@ -211,5 +211,5 @@ def test_plot_z_phase_calibration_result():
 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 = transform_circuit(circuit=c, target=cirq.CZ, replacement_map=replacement_map)
+    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)))

From 5a77eac72fb6db217b7c1fee6e0bd0d67cc1fa09 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Thu, 12 Dec 2024 15:07:56 -0800
Subject: [PATCH 04/10] fix replacement

---
 .../cirq/experiments/z_phase_calibration.py   | 43 +++++++++++++------
 .../experiments/z_phase_calibration_test.py   |  4 +-
 cirq-core/cirq/ops/fsim_gate.py               | 29 +++++++++++++
 3 files changed, 63 insertions(+), 13 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index d67f552932e..179f6d69cb7 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -23,7 +23,7 @@
 from cirq.experiments import xeb_fitting
 from cirq.experiments.two_qubit_xeb import parallel_xeb_workflow
 from cirq.transformers import transformer_api
-from cirq import ops, circuits
+from cirq import ops, circuits, protocols
 
 if TYPE_CHECKING:
     import cirq
@@ -279,20 +279,27 @@ class CalibrationTransformer:
 
     def __init__(
         self,
-        target: Union['cirq.Gate', 'cirq.GateFamily', 'cirq.Gateset'],
-        replacement_map: Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate'],
+        target: 'Cirq.Gate',
+        calibration_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:
+            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 = ops.GateFamily(target) if isinstance(target, ops.Gate) else target
-        self.replacement_map = replacement_map
+        self.target = target
+        self.target_as_fsim = (
+            target
+            if isinstance(target, ops.PhasedFSimGate)
+            else ops.PhasedFSimGate.from_matrix(protocols.unitary(target))
+        )
+        if self.target is None:
+            raise ValueError(f"{target} is not equivalent to a PhasedFSimGate")
+        self.calibration_map = calibration_map
 
     def __call__(
         self,
@@ -308,23 +315,35 @@ def __call__(
 
         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`.
+                in `calibration_map` replaced by the gate that cancels the effect of z-phases.
         """
         new_moments = []
         for moment in circuit:
             new_moment = []
             for op in moment:
-                if op not in self.target:
+                if op.gate != self.target:
                     # not a target.
                     new_moment.append(op)
                     continue
-                gate = self.replacement_map.get(op.qubits, None) or self.replacement_map.get(
+                gate = self.calibration_map.get(op.qubits, None) or self.calibration_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))
+                delta_theta = gate.theta - self.target_as_fsim.theta
+                delta_phi = gate.phi - self.target_as_fsim.phi
+                delta_chi = gate.chi - self.target_as_fsim.chi
+                delta_zeta = gate.zeta - self.target_as_fsim.zeta
+                delta_gamma = gate.gamma - self.target_as_fsim.gamma
+                new_gate = ops.PhasedFSimGate(
+                    theta=self.target_as_fsim.theta - delta_theta,
+                    gamma=self.target_as_fsim.gamma - delta_gamma,
+                    phi=self.target_as_fsim.phi - delta_phi,
+                    zeta=self.target_as_fsim.zeta - delta_zeta,
+                    chi=self.target_as_fsim.chi - delta_chi,
+                )
+                new_moment.append(new_gate(*op.qubits))
             new_moments.append(new_moment)
         return circuits.Circuit.from_moments(*new_moments)
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 8ad90c83b95..3f6ec5ba406 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -212,4 +212,6 @@ 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)))
+    assert new_circuit == cirq.Circuit(
+        cirq.PhasedFSimGate(0, -1, -2, -3, 2 * np.pi - 4).on(cirq.q(0), cirq.q(1))
+    )
diff --git a/cirq-core/cirq/ops/fsim_gate.py b/cirq-core/cirq/ops/fsim_gate.py
index e323d971f48..61117bfc90b 100644
--- a/cirq-core/cirq/ops/fsim_gate.py
+++ b/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."""

From 336ef693431c21adb634a06de478226e4911dd25 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Thu, 12 Dec 2024 15:37:07 -0800
Subject: [PATCH 05/10] add tests

---
 .../cirq/experiments/z_phase_calibration.py     |  1 +
 cirq-core/cirq/ops/fsim_gate_test.py            | 17 +++++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 179f6d69cb7..cb5829a639c 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -325,6 +325,7 @@ def __call__(
                     # not a target.
                     new_moment.append(op)
                     continue
+                assert len(op.qubits) == 2
                 gate = self.calibration_map.get(op.qubits, None) or self.calibration_map.get(
                     op.qubits[::-1], None
                 )
diff --git a/cirq-core/cirq/ops/fsim_gate_test.py b/cirq-core/cirq/ops/fsim_gate_test.py
index 684f36dffa8..c68eacdf2ce 100644
--- a/cirq-core/cirq/ops/fsim_gate_test.py
+++ b/cirq-core/cirq/ops/fsim_gate_test.py
@@ -797,3 +797,20 @@ 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)

From fe2e34784b7daf4e22787b548ddb68fa27fb3a56 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Thu, 12 Dec 2024 15:54:00 -0800
Subject: [PATCH 06/10] types

---
 cirq-core/cirq/experiments/z_phase_calibration.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 4a5c0b7ab06..e5a99a76b2b 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -304,12 +304,11 @@ def __init__(
                 calling `calibrate_z_phases`.
         """
         self.target = target
-        self.target_as_fsim = (
-            target
-            if isinstance(target, ops.PhasedFSimGate)
-            else ops.PhasedFSimGate.from_matrix(protocols.unitary(target))
-        )
-        if self.target is None:
+        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
 

From 2920112899aca8288b4c0d05691fcb75ef630db7 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Thu, 12 Dec 2024 15:55:08 -0800
Subject: [PATCH 07/10] typo

---
 cirq-core/cirq/experiments/z_phase_calibration.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index e5a99a76b2b..bdc6e1039de 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -291,7 +291,7 @@ class CalibrationTransformer:
 
     def __init__(
         self,
-        target: 'Cirq.Gate',
+        target: 'cirq.Gate',
         calibration_map: Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate'],
     ):
         """Create a CalibrationTransformer that replaces gates matching `target`.

From 95e3cfa704458a9b42b072e49a7ba576ff9badd0 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 13 Dec 2024 16:53:54 -0800
Subject: [PATCH 08/10] use ZPowGates

---
 .../cirq/experiments/z_phase_calibration.py   | 50 +++++++++++--------
 .../experiments/z_phase_calibration_test.py   | 23 ++++++---
 2 files changed, 45 insertions(+), 28 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index bdc6e1039de..2604c3776c9 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -286,6 +286,15 @@ def plot_z_phase_calibration_result(
     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:
 
@@ -294,7 +303,10 @@ def __init__(
         target: 'cirq.Gate',
         calibration_map: Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate'],
     ):
-        """Create a CalibrationTransformer that replaces gates matching `target`.
+        """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
@@ -318,44 +330,38 @@ def __call__(
         *,
         context: Optional[transformer_api.TransformerContext] = None,
     ) -> 'cirq.Circuit':
-        """Replace every occurance of a calibrated gate with a proper replacement.
+        """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 each gate matching `target` and acting on a pair
-                in `calibration_map` replaced by the gate that cancels the effect of z-phases.
+            New circuit with the extra ZPowGates.
         """
         new_moments = []
         for moment in circuit:
-            new_moment = []
+            before = []
+            after = []
             for op in moment:
                 if op.gate != self.target:
                     # not a target.
-                    new_moment.append(op)
                     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 calibrated version exist
-                    new_moment.append(op)
+                    # no calibration available.
                     continue
-                delta_theta = gate.theta - self.target_as_fsim.theta
-                delta_phi = gate.phi - self.target_as_fsim.phi
-                delta_chi = gate.chi - self.target_as_fsim.chi
-                delta_zeta = gate.zeta - self.target_as_fsim.zeta
-                delta_gamma = gate.gamma - self.target_as_fsim.gamma
-                new_gate = ops.PhasedFSimGate(
-                    theta=self.target_as_fsim.theta - delta_theta,
-                    gamma=self.target_as_fsim.gamma - delta_gamma,
-                    phi=self.target_as_fsim.phi - delta_phi,
-                    zeta=self.target_as_fsim.zeta - delta_zeta,
-                    chi=self.target_as_fsim.chi - delta_chi,
-                )
-                new_moment.append(new_gate(*op.qubits))
-            new_moments.append(new_moment)
+                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)
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 3f6ec5ba406..fd15064ed23 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -208,10 +208,21 @@ def test_plot_z_phase_calibration_result():
     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, 2 * np.pi - 4).on(cirq.q(0), cirq.q(1))
+@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))

From f49d940672ac3749528076cb37b36664e0b5c3ec Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 18 Dec 2024 14:06:42 -0800
Subject: [PATCH 09/10] fix types and coverage

---
 cirq-core/cirq/experiments/z_phase_calibration.py      |  4 ++--
 cirq-core/cirq/experiments/z_phase_calibration_test.py | 10 ++++++++++
 cirq-core/cirq/ops/fsim_gate_test.py                   |  4 ++++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 2604c3776c9..a4faef2c3a7 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/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
 
@@ -339,7 +339,7 @@ def __call__(
         Returns:
             New circuit with the extra ZPowGates.
         """
-        new_moments = []
+        new_moments: List[Union[List[cirq.Operation], 'cirq.Moment']] = []
         for moment in circuit:
             before = []
             after = []
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index fd15064ed23..c29bf54d5ba 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -226,3 +226,13 @@ def test_transform_circuit(angles):
         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)
diff --git a/cirq-core/cirq/ops/fsim_gate_test.py b/cirq-core/cirq/ops/fsim_gate_test.py
index c68eacdf2ce..a4d00c87310 100644
--- a/cirq-core/cirq/ops/fsim_gate_test.py
+++ b/cirq-core/cirq/ops/fsim_gate_test.py
@@ -814,3 +814,7 @@ def test_phased_fsim_json_dict():
 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

From eb66158cd667a3e58a0cb29dddd0a4b7202618e8 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 20 Dec 2024 11:38:29 -0800
Subject: [PATCH 10/10] address comments

---
 cirq-core/cirq/experiments/z_phase_calibration_test.py |  4 ++--
 cirq-core/cirq/ops/fsim_gate.py                        | 10 ++++++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index c29bf54d5ba..5f30c86843a 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -208,7 +208,7 @@ def test_plot_z_phase_calibration_result():
     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)))
+@pytest.mark.parametrize('angles', 2 * np.pi * np.random.random((10, 10)))
 def test_transform_circuit(angles):
     theta, phi = angles[:2]
     old_zs = angles[2:6]
@@ -225,7 +225,7 @@ def test_transform_circuit(angles):
     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))
+    np.testing.assert_allclose(cirq.unitary(circuit_with_replacement_gate), cirq.unitary(c))
 
 
 def test_transform_circuit_invalid_gate_raises():
diff --git a/cirq-core/cirq/ops/fsim_gate.py b/cirq-core/cirq/ops/fsim_gate.py
index 61117bfc90b..45cb3798bbf 100644
--- a/cirq-core/cirq/ops/fsim_gate.py
+++ b/cirq-core/cirq/ops/fsim_gate.py
@@ -349,6 +349,16 @@ def from_fsim_rz(
 
     @staticmethod
     def from_matrix(u: np.ndarray) -> Optional['PhasedFSimGate']:
+        """Contruct a PhasedFSimGate from unitary.
+
+        Args:
+            u: A unitary matrix representing a PhasedFSimGate.
+
+        Returns:
+            - Either PhasedFSimGate with the given unitary or None if
+                the matrix is not unitary or if doesn't represent a 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]