cirq-core 1.4.0.dev20240529225110__py3-none-any.whl → 1.5.0__py3-none-any.whl

cirq/experiments/z_phase_calibration.py

@@ -0,0 +1,368 @@
+# Copyright 2024 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Provides a method to do z-phase calibration for excitation-preserving gates."""
+import multiprocessing
+import multiprocessing.pool
+from typing import Any, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from cirq import circuits, ops, protocols
+from cirq.experiments import xeb_fitting
+from cirq.experiments.two_qubit_xeb import parallel_xeb_workflow
+from cirq.transformers import transformer_api
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+    import cirq
+
+
+def z_phase_calibration_workflow(
+    sampler: 'cirq.Sampler',
+    qubits: Optional[Sequence['cirq.GridQubit']] = None,
+    two_qubit_gate: 'cirq.Gate' = ops.CZ,
+    options: Optional[xeb_fitting.XEBPhasedFSimCharacterizationOptions] = None,
+    n_repetitions: int = 10**4,
+    n_combinations: int = 10,
+    n_circuits: int = 20,
+    cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
+    random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    atol: float = 1e-3,
+    num_workers_or_pool: Union[int, 'multiprocessing.pool.Pool'] = -1,
+    pairs: Optional[Sequence[Tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
+    tags: Sequence[Any] = (),
+) -> Tuple[xeb_fitting.XEBCharacterizationResult, 'pd.DataFrame']:
+    """Perform z-phase calibration for excitation-preserving gates.
+
+    For a given excitation-preserving two-qubit gate we assume an error model that can be described
+    using Z-rotations:
+                0: ───Rz(a)───two_qubit_gate───Rz(c)───
+                                │
+                1: ───Rz(b)───two_qubit_gate───Rz(d)───
+    for some angles a, b, c, and d.
+
+    Since the two-qubit gate is a excitation-preserving-gate, it can be represented by an FSimGate
+    and the effect of rotations turns it into a PhasedFSimGate. Using XEB-data we find the
+    PhasedFSimGate parameters that minimize the infidelity of the gate.
+
+    References:
+        - https://arxiv.org/abs/2001.08343
+        - https://arxiv.org/abs/2010.07965
+        - https://arxiv.org/abs/1910.11333
+
+    Args:
+        sampler: The quantum engine or simulator to run the circuits.
+        qubits: Qubits to use. If none, use all qubits on the sampler's device.
+        two_qubit_gate: The entangling gate to use.
+        options: The XEB-fitting options. If None, calibrate only the three phase angles
+            (chi, gamma, zeta) using the representation of a two-qubit gate as an FSimGate
+            for the initial guess.
+        n_repetitions: The number of repetitions to use.
+        n_combinations: The number of combinations to generate.
+        n_circuits: The number of circuits to generate.
+        cycle_depths: The cycle depths to use.
+        random_state: The random state to use.
+        atol: Absolute tolerance to be used by the minimizer.
+        num_workers_or_pool: An optional multi-processing pool or number of workers.
+            A zero value means no multiprocessing.
+            A positive integer value will create a pool with the given number of workers.
+            A negative value will create pool with maximum number of workers.
+        pairs: Pairs to use. If not specified, use all pairs between adjacent qubits.
+        tags: Tags to add to two qubit operations.
+    Returns:
+        - An `XEBCharacterizationResult` object that contains the calibration result.
+        - A `pd.DataFrame` comparing the before and after fidelities.
+    """
+
+    pool: Optional['multiprocessing.pool.Pool'] = None
+    local_pool = False
+    if isinstance(num_workers_or_pool, multiprocessing.pool.Pool):
+        pool = num_workers_or_pool  # pragma: no cover
+    elif num_workers_or_pool != 0:
+        pool = multiprocessing.Pool(num_workers_or_pool if num_workers_or_pool > 0 else None)
+        local_pool = True
+
+    fids_df_0, circuits, sampled_df = parallel_xeb_workflow(
+        sampler=sampler,
+        qubits=qubits,
+        entangling_gate=two_qubit_gate,
+        n_repetitions=n_repetitions,
+        cycle_depths=cycle_depths,
+        n_circuits=n_circuits,
+        n_combinations=n_combinations,
+        random_state=random_state,
+        pool=pool,
+        tags=tags,
+        pairs=pairs,
+    )
+
+    if options is None:
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
+            characterize_chi=True,
+            characterize_gamma=True,
+            characterize_zeta=True,
+            characterize_theta=False,
+            characterize_phi=False,
+        ).with_defaults_from_gate(two_qubit_gate)
+
+    p_circuits = [
+        xeb_fitting.parameterize_circuit(circuit, options, ops.GateFamily(two_qubit_gate))
+        for circuit in circuits
+    ]
+
+    result = xeb_fitting.characterize_phased_fsim_parameters_with_xeb_by_pair(
+        sampled_df=sampled_df,
+        parameterized_circuits=p_circuits,
+        cycle_depths=cycle_depths,
+        options=options,
+        fatol=atol,
+        xatol=atol,
+        pool=pool,
+    )
+
+    before_after = xeb_fitting.before_and_after_characterization(
+        fids_df_0, characterization_result=result
+    )
+
+    if local_pool:
+        assert isinstance(pool, multiprocessing.pool.Pool)
+        pool.close()
+    return result, before_after
+
+
+def calibrate_z_phases(
+    sampler: 'cirq.Sampler',
+    qubits: Optional[Sequence['cirq.GridQubit']] = None,
+    two_qubit_gate: 'cirq.Gate' = ops.CZ,
+    options: Optional[xeb_fitting.XEBPhasedFSimCharacterizationOptions] = None,
+    n_repetitions: int = 10**4,
+    n_combinations: int = 10,
+    n_circuits: int = 20,
+    cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
+    random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    atol: float = 1e-3,
+    num_workers_or_pool: Union[int, 'multiprocessing.pool.Pool'] = -1,
+    pairs: Optional[Sequence[Tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
+    tags: Sequence[Any] = (),
+) -> Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate']:
+    """Perform z-phase calibration for excitation-preserving gates.
+
+    For a given excitation-preserving two-qubit gate we assume an error model that can be described
+    using Z-rotations:
+                0: ───Rz(a)───two_qubit_gate───Rz(c)───
+                                │
+                1: ───Rz(b)───two_qubit_gate───Rz(d)───
+    for some angles a, b, c, and d.
+
+    Since the two-qubit gate is a excitation-preserving gate, it can be represented by an FSimGate
+    and the effect of rotations turns it into a PhasedFSimGate. Using XEB-data we find the
+    PhasedFSimGate parameters that minimize the infidelity of the gate.
+
+    References:
+        - https://arxiv.org/abs/2001.08343
+        - https://arxiv.org/abs/2010.07965
+        - https://arxiv.org/abs/1910.11333
+
+    Args:
+        sampler: The quantum engine or simulator to run the circuits.
+        qubits: Qubits to use. If none, use all qubits on the sampler's device.
+        two_qubit_gate: The entangling gate to use.
+        options: The XEB-fitting options. If None, calibrate only the three phase angles
+            (chi, gamma, zeta) using the representation of a two-qubit gate as an FSimGate
+            for the initial guess.
+        n_repetitions: The number of repetitions to use.
+        n_combinations: The number of combinations to generate.
+        n_circuits: The number of circuits to generate.
+        cycle_depths: The cycle depths to use.
+        random_state: The random state to use.
+        atol: Absolute tolerance to be used by the minimizer.
+        num_workers_or_pool: An optional multi-processing pool or number of workers.
+            A zero value means no multiprocessing.
+            A positive integer value will create a pool with the given number of workers.
+            A negative value will create pool with maximum number of workers.
+        pairs: Pairs to use. If not specified, use all pairs between adjacent qubits.
+        tags: Tags to add to two qubit operations.
+
+    Returns:
+        - A dictionary mapping qubit pairs to the calibrated PhasedFSimGates.
+    """
+
+    if options is None:
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
+            characterize_chi=True,
+            characterize_gamma=True,
+            characterize_zeta=True,
+            characterize_theta=False,
+            characterize_phi=False,
+        ).with_defaults_from_gate(two_qubit_gate)
+
+    result, _ = z_phase_calibration_workflow(
+        sampler=sampler,
+        qubits=qubits,
+        two_qubit_gate=two_qubit_gate,
+        options=options,
+        n_repetitions=n_repetitions,
+        n_combinations=n_combinations,
+        n_circuits=n_circuits,
+        cycle_depths=cycle_depths,
+        random_state=random_state,
+        atol=atol,
+        num_workers_or_pool=num_workers_or_pool,
+        tags=tags,
+        pairs=pairs,
+    )
+
+    gates = {}
+    for pair, params in result.final_params.items():
+        params['theta'] = params.get('theta', options.theta_default or 0)
+        params['phi'] = params.get('phi', options.phi_default or 0)
+        params['zeta'] = params.get('zeta', options.zeta_default or 0)
+        params['chi'] = params.get('chi', options.chi_default or 0)
+        params['gamma'] = params.get('gamma', options.gamma_default or 0)
+        gates[pair] = ops.PhasedFSimGate(**params)
+    return gates
+
+
+def plot_z_phase_calibration_result(
+    before_after_df: 'pd.DataFrame',
+    axes: Optional[np.ndarray[Sequence[Sequence['plt.Axes']], np.dtype[np.object_]]] = None,
+    pairs: Optional[Sequence[Tuple['cirq.Qid', 'cirq.Qid']]] = None,
+    *,
+    with_error_bars: bool = False,
+) -> np.ndarray[Sequence[Sequence['plt.Axes']], np.dtype[np.object_]]:
+    """A helper method to plot the result of running z-phase calibration.
+
+    Note that the plotted fidelity is a statistical estimate of the true fidelity and as a result
+    may be outside the [0, 1] range.
+
+    Args:
+        before_after_df: The second return object of running `z_phase_calibration_workflow`.
+        axes: And ndarray of the axes to plot on.
+            The number of axes is expected to be >= number of qubit pairs.
+        pairs: If provided, only the given pairs are plotted.
+        with_error_bars: Whether to add error bars or not.
+            The width of the bar is an upper bound on standard variation of the estimated fidelity.
+    """
+    if pairs is None:
+        pairs = before_after_df.index
+    if axes is None:
+        # Create a 16x9 rectangle.
+        ncols = int(np.ceil(np.sqrt(9 / 16 * len(pairs))))
+        nrows = (len(pairs) + ncols - 1) // ncols
+        _, axes = plt.subplots(nrows=nrows, ncols=ncols)
+    axes = axes if isinstance(axes, np.ndarray) else np.array(axes)
+    for pair, ax in zip(pairs, axes.flatten()):
+        row = before_after_df.loc[[pair]].iloc[0]
+        ax.errorbar(
+            row.cycle_depths_0,
+            row.fidelities_0,
+            yerr=row.layer_fid_std_0 * with_error_bars,
+            label='original',
+        )
+        ax.errorbar(
+            row.cycle_depths_0,
+            row.fidelities_c,
+            yerr=row.layer_fid_std_c * with_error_bars,
+            label='calibrated',
+        )
+        ax.axhline(1, linestyle='--')
+        ax.set_xlabel('cycle depth')
+        ax.set_ylabel('fidelity estimate')
+        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/experiments/z_phase_calibration_test.py

@@ -0,0 +1,241 @@
+# Copyright 2024 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pandas as pd
+import pytest
+
+import cirq
+from cirq.experiments.xeb_fitting import XEBPhasedFSimCharacterizationOptions
+from cirq.experiments.z_phase_calibration import (
+    calibrate_z_phases,
+    CalibrationTransformer,
+    plot_z_phase_calibration_result,
+    z_phase_calibration_workflow,
+)
+
+_ANGLES = ['theta', 'phi', 'chi', 'zeta', 'gamma']
+# fix random generator seed to ensure reproducibility and faster convergence
+_SEED = 276154030
+
+
+def _create_tests(n, with_options: bool = False):
+    rng = np.random.default_rng(_SEED)
+    angles = (rng.random((n, 5)) * 2 - 1) * np.pi
+    # Add errors to the last 3 angles (chi, zeta, gamma).
+    # The errors are in the union (-2, -1) U (1, 2).
+    # This is because we run the tests with few repetitions so a small error might not get fixed.
+    error = np.concatenate(
+        [np.zeros((n, 2)), (rng.random((n, 3)) + 1) * rng.choice([-1, 1], (n, 3))], axis=-1
+    )
+    if with_options:
+        options = []
+        for _ in range(n):
+            v = [False, False, False]
+            # Calibrate only one to keep the run time down.
+            v[rng.integers(0, 3)] = True
+            options.append(
+                {
+                    'characterize_chi': v[0],
+                    'characterize_gamma': v[1],
+                    'characterize_zeta': v[2],
+                    'characterize_phi': False,
+                    'characterize_theta': False,
+                }
+            )
+
+        return zip(angles, error, options)
+    return zip(angles, error)
+
+
+def _trace_distance(A, B):
+    return 0.5 * np.abs(np.linalg.eigvals(A - B)).sum()
+
+
+class _TestSimulator(cirq.Simulator):
+    """A simulator that replaces a specific gate by another."""
+
+    def __init__(self, gate: cirq.Gate, replacement: cirq.Gate, **kwargs):
+        super().__init__(**kwargs)
+        self.gate = gate
+        self.replacement = replacement
+
+    def _core_iterator(
+        self,
+        circuit: 'cirq.AbstractCircuit',
+        sim_state,
+        all_measurements_are_terminal: bool = False,
+    ):
+        new_circuit = cirq.Circuit(
+            [
+                [op if op.gate != self.gate else self.replacement(*op.qubits) for op in m]
+                for m in circuit
+            ]
+        )
+        yield from super()._core_iterator(new_circuit, sim_state, all_measurements_are_terminal)
+
+
+@pytest.mark.parametrize(
+    ['angles', 'error', 'characterization_flags'], _create_tests(n=10, with_options=True)
+)
+def test_calibrate_z_phases(angles, error, characterization_flags):
+
+    original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})
+    actual_gate = cirq.PhasedFSimGate(**{k: v + e for k, v, e in zip(_ANGLES, angles, error)})
+
+    options = XEBPhasedFSimCharacterizationOptions(
+        **{f'{n}_default': t for n, t in zip(_ANGLES, angles)}, **characterization_flags
+    )
+
+    sampler = _TestSimulator(original_gate, actual_gate, seed=_SEED)
+    qubits = cirq.q(0, 0), cirq.q(0, 1)
+    calibrated_gate = calibrate_z_phases(
+        sampler,
+        qubits,
+        original_gate,
+        options,
+        n_repetitions=10,
+        n_combinations=10,
+        n_circuits=10,
+        cycle_depths=range(3, 10),
+        random_state=_SEED,
+    )[qubits]
+
+    initial_unitary = cirq.unitary(original_gate)
+    final_unitary = cirq.unitary(calibrated_gate)
+    target_unitary = cirq.unitary(actual_gate)
+    maximally_mixed_state = np.eye(4) / 2
+    dm_initial = initial_unitary @ maximally_mixed_state @ initial_unitary.T.conj()
+    dm_final = final_unitary @ maximally_mixed_state @ final_unitary.T.conj()
+    dm_target = target_unitary @ maximally_mixed_state @ target_unitary.T.conj()
+
+    original_dist = _trace_distance(dm_initial, dm_target)
+    new_dist = _trace_distance(dm_final, dm_target)
+
+    # Either we reduced the error or the error is small enough.
+    assert new_dist < original_dist or new_dist < 1e-6
+
+
+@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=3))
+def test_calibrate_z_phases_no_options(angles, error):
+
+    original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})
+    actual_gate = cirq.PhasedFSimGate(**{k: v + e for k, v, e in zip(_ANGLES, angles, error)})
+
+    sampler = _TestSimulator(original_gate, actual_gate, seed=_SEED)
+    qubits = cirq.q(0, 0), cirq.q(0, 1)
+    calibrated_gate = calibrate_z_phases(
+        sampler,
+        qubits,
+        original_gate,
+        options=None,
+        n_repetitions=10,
+        n_combinations=10,
+        n_circuits=10,
+        cycle_depths=range(3, 10),
+        random_state=_SEED,
+    )[qubits]
+
+    initial_unitary = cirq.unitary(original_gate)
+    final_unitary = cirq.unitary(calibrated_gate)
+    target_unitary = cirq.unitary(actual_gate)
+    maximally_mixed_state = np.eye(4) / 2
+    dm_initial = initial_unitary @ maximally_mixed_state @ initial_unitary.T.conj()
+    dm_final = final_unitary @ maximally_mixed_state @ final_unitary.T.conj()
+    dm_target = target_unitary @ maximally_mixed_state @ target_unitary.T.conj()
+
+    original_dist = _trace_distance(dm_initial, dm_target)
+    new_dist = _trace_distance(dm_final, dm_target)
+
+    # Either we reduced the error or the error is small enough.
+    assert new_dist < original_dist or new_dist < 1e-6
+
+
+@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=3))
+def test_calibrate_z_phases_workflow_no_options(angles, error):
+
+    original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})
+    actual_gate = cirq.PhasedFSimGate(**{k: v + e for k, v, e in zip(_ANGLES, angles, error)})
+
+    sampler = _TestSimulator(original_gate, actual_gate, seed=_SEED)
+    qubits = cirq.q(0, 0), cirq.q(0, 1)
+    result, _ = z_phase_calibration_workflow(
+        sampler,
+        qubits,
+        original_gate,
+        options=None,
+        n_repetitions=1,
+        n_combinations=1,
+        n_circuits=1,
+        cycle_depths=(1, 2),
+        random_state=_SEED,
+    )
+
+    for params in result.final_params.values():
+        assert 'zeta' in params
+        assert 'chi' in params
+        assert 'gamma' in params
+        assert 'phi' not in params
+        assert 'theta' not in params
+
+
+def test_plot_z_phase_calibration_result():
+    df = pd.DataFrame()
+    qs = cirq.q(0, 0), cirq.q(0, 1), cirq.q(0, 2)
+    df.index = [qs[:2], qs[-2:]]
+    df['cycle_depths_0'] = [[1, 2, 3]] * 2
+    df['fidelities_0'] = [[0.9, 0.8, 0.7], [0.6, 0.4, 0.1]]
+    df['layer_fid_std_0'] = [0.1, 0.2]
+    df['fidelities_c'] = [[0.9, 0.92, 0.93], [0.7, 0.77, 0.8]]
+    df['layer_fid_std_c'] = [0.2, 0.3]
+
+    axes = plot_z_phase_calibration_result(before_after_df=df)
+
+    np.testing.assert_allclose(axes[0].lines[0].get_xdata().astype(float), [1, 2, 3])
+    np.testing.assert_allclose(axes[0].lines[0].get_ydata().astype(float), [0.9, 0.8, 0.7])
+    np.testing.assert_allclose(axes[0].lines[1].get_ydata().astype(float), [0.9, 0.92, 0.93])
+
+    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', 2 * np.pi * 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(c))
+
+
+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/interop/__init__.py

@@ -13,4 +13,7 @@
 # limitations under the License.
 """Methods for interoperating with other quantum software."""
 
-from cirq.interop.quirk import quirk_json_to_circuit, quirk_url_to_circuit
+from cirq.interop.quirk import (
+    quirk_json_to_circuit as quirk_json_to_circuit,
+    quirk_url_to_circuit as quirk_url_to_circuit,
+)

cirq/interop/quirk/__init__.py

@@ -20,9 +20,12 @@ References:
 
 # Imports from cells are only to ensure operation reprs work correctly.
 from cirq.interop.quirk.cells import (
-    QuirkArithmeticGate,
-    QuirkInputRotationOperation,
-    QuirkQubitPermutationGate,
+    QuirkArithmeticGate as QuirkArithmeticGate,
+    QuirkInputRotationOperation as QuirkInputRotationOperation,
+    QuirkQubitPermutationGate as QuirkQubitPermutationGate,
 )
 
-from cirq.interop.quirk.url_to_circuit import quirk_json_to_circuit, quirk_url_to_circuit
+from cirq.interop.quirk.url_to_circuit import (
+    quirk_json_to_circuit as quirk_json_to_circuit,
+    quirk_url_to_circuit as quirk_url_to_circuit,
+)

cirq/interop/quirk/cells/__init__.py

@@ -13,17 +13,28 @@
 # limitations under the License.
 """This module defines building blocks for parsing Quirk circuits."""
 
-from cirq.interop.quirk.cells.all_cells import generate_all_quirk_cell_makers
+from cirq.interop.quirk.cells.all_cells import (
+    generate_all_quirk_cell_makers as generate_all_quirk_cell_makers,
+)
 
-from cirq.interop.quirk.cells.cell import Cell, CellMaker, CellMakerArgs, ExplicitOperationsCell
+from cirq.interop.quirk.cells.cell import (
+    Cell as Cell,
+    CellMaker as CellMaker,
+    CellMakerArgs as CellMakerArgs,
+    ExplicitOperationsCell as ExplicitOperationsCell,
+)
 
-from cirq.interop.quirk.cells.composite_cell import CompositeCell
+from cirq.interop.quirk.cells.composite_cell import CompositeCell as CompositeCell
 
-from cirq.interop.quirk.cells.qubit_permutation_cells import QuirkQubitPermutationGate
+from cirq.interop.quirk.cells.qubit_permutation_cells import (
+    QuirkQubitPermutationGate as QuirkQubitPermutationGate,
+)
 
-from cirq.interop.quirk.cells.arithmetic_cells import QuirkArithmeticGate
+from cirq.interop.quirk.cells.arithmetic_cells import QuirkArithmeticGate as QuirkArithmeticGate
 
-from cirq.interop.quirk.cells.input_rotation_cells import QuirkInputRotationOperation
+from cirq.interop.quirk.cells.input_rotation_cells import (
+    QuirkInputRotationOperation as QuirkInputRotationOperation,
+)
 
 import cirq.interop.quirk.cells.swap_cell
 import cirq.interop.quirk.cells.control_cells