Expand Floquet calibration to arbitrary FSim gates (#4164)

* Added a method try_convert_gate_to_fsim which can convert to PhaseCalibratedFSimGate wide range of gates, including sqrt(iSWAP), Sycamore and Controlled-Z gates.
* Used this method (instead of try_convert_sqrt_iswap_to_fsim) as default gate translator in all methods in workflow.py.

As a result, compensation step Floquet calibration now works not only for sqrt(iSWAP), but also for:
* all FSim gates (including Sycamore gate);
* PhasedFSim gates (such that zeta=gamma=0); 
* all ISwapPow gates;
* all PhasedISwapPow gates;
* CZPow gates without global phase (including CZ).

This change doesn't affect characterization step the Floquet calibration.

cirq-google/cirq_google/calibration/__init__.py

@@ -34,6 +34,7 @@
     THETA_ZETA_GAMMA_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     merge_matching_results,
+    try_convert_sqrt_iswap_to_fsim,
 )
 
 from cirq_google.calibration.workflow import (
@@ -50,5 +51,4 @@
     run_calibrations,
     run_floquet_characterization_for_moments,
     run_zeta_chi_gamma_compensation_for_moments,
-    try_convert_sqrt_iswap_to_fsim,
 )

cirq-google/cirq_google/calibration/engine_simulator.py

@@ -23,6 +23,7 @@
     PhasedFSimCalibrationResult,
     PhasedFSimCharacterization,
     SQRT_ISWAP_INV_PARAMETERS,
+    try_convert_gate_to_fsim,
     try_convert_sqrt_iswap_to_fsim,
 )
 
@@ -244,6 +245,60 @@ def sample_gate(
             simulator, drift_generator=sample_gate, gates_translator=try_convert_sqrt_iswap_to_fsim
         )
 
+    @classmethod
+    def create_from_dictionary(
+        cls,
+        parameters: Dict[
+            Tuple[cirq.Qid, cirq.Qid], Dict[cirq.FSimGate, Union[PhasedFSimCharacterization, Dict]]
+        ],
+        *,
+        simulator: Optional[cirq.Simulator] = None,
+    ) -> 'PhasedFSimEngineSimulator':
+        """Creates PhasedFSimEngineSimulator with fixed drifts.
+
+        Args:
+            parameters: maps every pair of qubits and engine gate on that pair to a
+                characterization for that gate.
+            simulator: Simulator object to use. When None, a new instance of cirq.Simulator() will
+                be created.
+
+        Returns:
+            New PhasedFSimEngineSimulator instance.
+        """
+
+        for a, b in parameters.keys():
+            if a > b:
+                raise ValueError(
+                    f'All qubit pairs must be given in canonical order where the first qubit is '
+                    f'less than the second, got {a} > {b}'
+                )
+
+        def sample_gate(
+            a: cirq.Qid, b: cirq.Qid, gate: cirq.FSimGate
+        ) -> PhasedFSimCharacterization:
+            pair_parameters = None
+            swapped = False
+            if (a, b) in parameters:
+                pair_parameters = parameters[(a, b)].get(gate)
+            elif (b, a) in parameters:
+                pair_parameters = parameters[(b, a)].get(gate)
+                swapped = True
+
+            if pair_parameters is None:
+                raise ValueError(f'Missing parameters for value for pair {(a, b)} and gate {gate}.')
+            if not isinstance(pair_parameters, PhasedFSimCharacterization):
+                pair_parameters = PhasedFSimCharacterization(**pair_parameters)
+            if swapped:
+                pair_parameters = pair_parameters.parameters_for_qubits_swapped()
+
+            return pair_parameters
+
+        if simulator is None:
+            simulator = cirq.Simulator()
+        return cls(
+            simulator, drift_generator=sample_gate, gates_translator=try_convert_gate_to_fsim
+        )
+
     @classmethod
     def create_from_characterizations_sqrt_iswap(
         cls,

cirq-google/cirq_google/calibration/engine_simulator_test.py

@@ -22,6 +22,8 @@
 )
 import cirq
 
+SQRT_ISWAP_INV_GATE = cirq.FSimGate(np.pi / 4, 0.0)
+
 
 class DummyPhasedFSimCalibrationRequest(PhasedFSimCalibrationRequest):
     def to_calibration_layer(self) -> cirq_google.CalibrationLayer:
@@ -278,6 +280,50 @@ def test_from_dictionary_sqrt_iswap_simulates_correctly():
     assert cirq.allclose_up_to_global_phase(actual, expected)
 
 
+def test_create_from_dictionary_simulates_correctly():
+    parameters_ab_1 = {'theta': 0.6, 'zeta': 0.5, 'chi': 0.4, 'gamma': 0.3, 'phi': 0.2}
+    parameters_ab_2 = {'theta': 0.1, 'zeta': 0.2, 'chi': 0.3, 'gamma': 0.4, 'phi': 0.5}
+    parameters_bc = {'theta': 0.8, 'zeta': -0.5, 'chi': -0.4, 'gamma': -0.3, 'phi': -0.2}
+    parameters_cd = {'theta': 0.1, 'zeta': 0.2, 'chi': 0.3, 'gamma': 0.4, 'phi': 0.5}
+
+    a, b, c, d = cirq.LineQubit.range(4)
+    circuit = cirq.Circuit(
+        [
+            [cirq.X(a), cirq.Y(b), cirq.Z(c), cirq.H(d)],
+            [SQRT_ISWAP_INV_GATE.on(a, b), SQRT_ISWAP_INV_GATE.on(d, c)],
+            [SQRT_ISWAP_INV_GATE.on(b, c)],
+            [cirq_google.SYC.on(a, b), SQRT_ISWAP_INV_GATE.on(c, d)],
+        ]
+    )
+    expected_circuit = cirq.Circuit(
+        [
+            [cirq.X(a), cirq.Y(b), cirq.Z(c), cirq.H(d)],
+            [
+                cirq.PhasedFSimGate(**parameters_ab_1).on(a, b),
+                cirq.PhasedFSimGate(**parameters_cd).on(c, d),
+            ],
+            [cirq.PhasedFSimGate(**parameters_bc).on(b, c)],
+            [
+                cirq.PhasedFSimGate(**parameters_ab_2).on(a, b),
+                cirq.PhasedFSimGate(**parameters_cd).on(c, d),
+            ],
+        ]
+    )
+
+    engine_simulator = PhasedFSimEngineSimulator.create_from_dictionary(
+        parameters={
+            (a, b): {SQRT_ISWAP_INV_GATE: parameters_ab_1, cirq_google.SYC: parameters_ab_2},
+            (b, c): {SQRT_ISWAP_INV_GATE: parameters_bc},
+            (c, d): {SQRT_ISWAP_INV_GATE: parameters_cd},
+        }
+    )
+
+    actual = engine_simulator.final_state_vector(circuit)
+    expected = cirq.final_state_vector(expected_circuit)
+
+    assert cirq.allclose_up_to_global_phase(actual, expected)
+
+
 def test_from_dictionary_sqrt_iswap_ideal_when_missing_gate_fails():
     a, b = cirq.LineQubit.range(2)
     circuit = cirq.Circuit(cirq.FSimGate(np.pi / 4, 0.0).on(a, b))
@@ -444,6 +490,20 @@ def test_from_characterizations_sqrt_iswap_when_invalid_arguments_fails():
         )
 
 
+def test_create_from_dictionary_imvalid_parameters_fails():
+    a, b = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(cirq.CZ(a, b))
+
+    simulator = PhasedFSimEngineSimulator.create_from_dictionary({})
+    with pytest.raises(ValueError, match='Missing parameters'):
+        simulator.final_state_vector(circuit)
+
+    with pytest.raises(ValueError, match='canonical order'):
+        PhasedFSimEngineSimulator.create_from_dictionary(
+            parameters={(b, a): {'theta': 0.6, 'phi': 0.2}}
+        )
+
+
 def _create_sqrt_iswap_request(
     pairs: Iterable[Tuple[cirq.Qid, cirq.Qid]],
     options: FloquetPhasedFSimCalibrationOptions = ALL_ANGLES_FLOQUET_PHASED_FSIM_CHARACTERIZATION,

cirq-google/cirq_google/calibration/phased_fsim.py

@@ -41,6 +41,7 @@
 )
 from cirq_google.api import v2
 from cirq_google.engine import Calibration, CalibrationLayer, CalibrationResult, Engine, EngineJob
+from cirq_google.ops import SycamoreGate
 
 if TYPE_CHECKING:
     import cirq_google
@@ -996,6 +997,13 @@ def with_zeta_chi_gamma_compensated(
             (cirq.rz(0.5 * gamma - beta).on(a), cirq.rz(0.5 * gamma + beta).on(b)),
         )
 
+    def _unitary_(self) -> np.array:
+        """Implements Cirq's `unitary` protocol for this object."""
+        p = np.exp(-np.pi * 1j * self.phase_exponent)
+        return (
+            np.diag([1, p, 1 / p, 1]) @ cirq.unitary(self.engine_gate) @ np.diag([1, 1 / p, p, 1])
+        )
+
 
 def try_convert_sqrt_iswap_to_fsim(gate: cirq.Gate) -> Optional[PhaseCalibratedFSimGate]:
     """Converts an equivalent gate to FSimGate(theta=π/4, phi=0) if possible.
@@ -1038,3 +1046,54 @@ def try_convert_sqrt_iswap_to_fsim(gate: cirq.Gate) -> Optional[PhaseCalibratedF
         return PhaseCalibratedFSimGate(cirq.FSimGate(theta=np.pi / 4, phi=0.0), 0.5)
 
     return None
+
+
+def try_convert_gate_to_fsim(gate: cirq.Gate) -> Optional[PhaseCalibratedFSimGate]:
+    """Converts a gate to equivalent PhaseCalibratedFSimGate if possible.
+
+    Args:
+        gate: Gate to convert.
+
+    Returns:
+        If provided gate is equivalent to some PhaseCalibratedFSimGate, returns that gate.
+        Otherwise returns None.
+    """
+    if cirq.is_parameterized(gate):
+        return None
+
+    phi = 0.0
+    theta = 0.0
+    phase_exponent = 0.0
+    if isinstance(gate, SycamoreGate):
+        phi = np.pi / 6
+        theta = np.pi / 2
+    elif isinstance(gate, cirq.FSimGate):
+        theta = gate.theta
+        phi = gate.phi
+    elif isinstance(gate, cirq.ISwapPowGate):
+        if not np.isclose(np.exp(np.pi * 1j * gate.global_shift * gate.exponent), 1.0):
+            return None
+        theta = -gate.exponent * np.pi / 2
+    elif isinstance(gate, cirq.PhasedFSimGate):
+        if not np.isclose(gate.zeta, 0.0) or not np.isclose(gate.gamma, 0.0):
+            return None
+        theta = gate.theta
+        phi = gate.phi
+        phase_exponent = -gate.chi / (2 * np.pi)
+    elif isinstance(gate, cirq.PhasedISwapPowGate):
+        theta = -gate.exponent * np.pi / 2
+        phase_exponent = -gate.phase_exponent
+    elif isinstance(gate, cirq.ops.CZPowGate):
+        if not np.isclose(np.exp(np.pi * 1j * gate.global_shift * gate.exponent), 1.0):
+            return None
+        phi = -np.pi * gate.exponent
+    else:
+        return None
+
+    phi = phi % (2 * np.pi)
+    theta = theta % (2 * np.pi)
+    if theta >= np.pi:
+        theta = 2 * np.pi - theta
+        phase_exponent = phase_exponent + 0.5
+    phase_exponent %= 1
+    return PhaseCalibratedFSimGate(cirq.FSimGate(theta=theta, phi=phi), phase_exponent)

cirq-google/cirq_google/calibration/phased_fsim_test.py

@@ -18,6 +18,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import sympy
 from google.protobuf import text_format
 
 import cirq
@@ -35,6 +36,7 @@
     PhasedFSimCalibrationResult,
     WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     merge_matching_results,
+    try_convert_gate_to_fsim,
     try_convert_sqrt_iswap_to_fsim,
     XEBPhasedFSimCalibrationRequest,
     XEBPhasedFSimCalibrationOptions,
@@ -874,3 +876,95 @@ def test_options_phase_corrected_override():
         ).zeta_chi_gamma_correction_override()
         == PhasedFSimCharacterization()
     )
+
+
+def test_try_convert_gate_to_fsim():
+    def check(gate: cirq.Gate, expected: PhaseCalibratedFSimGate):
+        assert np.allclose(cirq.unitary(gate), cirq.unitary(expected))
+        assert try_convert_gate_to_fsim(gate) == expected
+
+    check(
+        cirq.FSimGate(theta=0.3, phi=0.5),
+        PhaseCalibratedFSimGate(cirq.FSimGate(theta=0.3, phi=0.5), 0.0),
+    )
+
+    check(
+        cirq.FSimGate(7 * np.pi / 4, 0.0),
+        PhaseCalibratedFSimGate(cirq.FSimGate(np.pi / 4, 0.0), 0.5),
+    )
+
+    check(
+        cirq.ISwapPowGate(exponent=-0.5),
+        PhaseCalibratedFSimGate(cirq.FSimGate(theta=0.25 * np.pi, phi=0.0), 0.0),
+    )
+
+    check(
+        cirq.ops.ISwapPowGate(exponent=0.8, global_shift=2.5),
+        PhaseCalibratedFSimGate(cirq.FSimGate(theta=0.4 * np.pi, phi=0.0), 0.5),
+    )
+
+    gate = cirq.ops.ISwapPowGate(exponent=0.3, global_shift=0.4)
+    assert try_convert_gate_to_fsim(gate) is None
+
+    check(
+        cirq.PhasedFSimGate(theta=0.2, phi=0.5, chi=1.5 * np.pi),
+        PhaseCalibratedFSimGate(cirq.FSimGate(theta=0.2, phi=0.5), 0.25),
+    )
+
+    gate = cirq.PhasedFSimGate(theta=0.2, phi=0.5, zeta=1.5 * np.pi)
+    assert try_convert_gate_to_fsim(gate) is None
+
+    check(
+        cirq.PhasedISwapPowGate(exponent=-0.5, phase_exponent=0.75),
+        PhaseCalibratedFSimGate(cirq.FSimGate(theta=0.25 * np.pi, phi=0.0), 0.25),
+    )
+
+    check(cirq.CZ, PhaseCalibratedFSimGate(cirq.FSimGate(theta=0.0, phi=np.pi), 0.0))
+
+    check(
+        cirq.ops.CZPowGate(exponent=0.3),
+        PhaseCalibratedFSimGate(cirq.FSimGate(theta=0.0, phi=-0.3 * np.pi), 0.0),
+    )
+
+    check(
+        cirq.ops.CZPowGate(exponent=0.8, global_shift=2.5),
+        PhaseCalibratedFSimGate(cirq.FSimGate(theta=0.0, phi=-0.8 * np.pi), 0.0),
+    )
+
+    gate = cirq.ops.CZPowGate(exponent=0.3, global_shift=0.4)
+    assert try_convert_gate_to_fsim(gate) is None
+
+    check(
+        cirq_google.ops.SYC,
+        PhaseCalibratedFSimGate(cirq.FSimGate(phi=np.pi / 6, theta=np.pi / 2), 0.0),
+    )
+
+    assert try_convert_gate_to_fsim(cirq.CX) is None
+
+    # Parameterized gates are not supported.
+    x = sympy.Symbol('x')
+    assert try_convert_gate_to_fsim(cirq.ops.ISwapPowGate(exponent=x)) is None
+    assert try_convert_gate_to_fsim(cirq.PhasedFSimGate(theta=x)) is None
+    assert try_convert_gate_to_fsim(cirq.PhasedISwapPowGate(exponent=x)) is None
+    assert try_convert_gate_to_fsim(cirq.CZPowGate(exponent=x)) is None
+
+
+# Test that try_convert_gate_to_fsim is extension of try_convert_sqrt_iswap_to_fsim.
+# In other words, that both function return the same result for all gates on which
+# try_convert_sqrt_iswap_to_fsim is defined.
+# TODO: instead of having multiple gate translators, we should have one, the most general, and
+# restrict it to different gate sets.
+def test_gate_translators_are_consistent():
+    def check(gate):
+        result1 = try_convert_gate_to_fsim(gate)
+        result2 = try_convert_sqrt_iswap_to_fsim(gate)
+        assert result1 == result2
+        assert result1 is not None
+
+    check(cirq.FSimGate(theta=np.pi / 4, phi=0))
+    check(cirq.FSimGate(theta=-np.pi / 4, phi=0))
+    check(cirq.FSimGate(theta=7 * np.pi / 4, phi=0))
+    check(cirq.PhasedFSimGate(theta=np.pi / 4, phi=0))
+    check(cirq.ISwapPowGate(exponent=0.5))
+    check(cirq.ISwapPowGate(exponent=-0.5))
+    check(cirq.PhasedISwapPowGate(exponent=0.5, phase_exponent=-0.5))

cirq-google/cirq_google/calibration/workflow.py

@@ -40,6 +40,7 @@
     WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     THETA_ZETA_GAMMA_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     merge_matching_results,
+    try_convert_gate_to_fsim,
     try_convert_sqrt_iswap_to_fsim,
     PhasedFSimCalibrationOptions,
     RequestT,
@@ -847,7 +848,7 @@ def make_zeta_chi_gamma_compensation_for_moments(
     *,
     gates_translator: Callable[
         [cirq.Gate], Optional[PhaseCalibratedFSimGate]
-    ] = try_convert_sqrt_iswap_to_fsim,
+    ] = try_convert_gate_to_fsim,
     merge_subsets: bool = True,
 ) -> CircuitWithCalibration:
     """Compensates circuit moments against errors in zeta, chi and gamma angles.
@@ -897,7 +898,7 @@ def make_zeta_chi_gamma_compensation_for_operations(
     characterizations: List[PhasedFSimCalibrationResult],
     gates_translator: Callable[
         [cirq.Gate], Optional[PhaseCalibratedFSimGate]
-    ] = try_convert_sqrt_iswap_to_fsim,
+    ] = try_convert_gate_to_fsim,
     permit_mixed_moments: bool = False,
 ) -> cirq.Circuit:
     """Compensates circuit operations against errors in zeta, chi and gamma angles.
@@ -951,7 +952,6 @@ def _make_zeta_chi_gamma_compensation(
     gates_translator: Callable[[cirq.Gate], Optional[PhaseCalibratedFSimGate]],
     permit_mixed_moments: bool,
 ) -> CircuitWithCalibration:
-
     if permit_mixed_moments:
         raise NotImplementedError('Mixed moments compensation ist supported yet')
 
@@ -992,6 +992,12 @@ def _make_zeta_chi_gamma_compensation(
             if parameters is None:
                 raise ValueError(f'Missing characterization data for moment {moment}')
 
+            if translated.engine_gate != parameters.gate:
+                raise ValueError(
+                    f"Engine gate {translated.engine_gate} doesn't match characterized gate "
+                    f'{parameters.gate}'
+                )
+
             pair_parameters = parameters.get_parameters(a, b)
             if pair_parameters is None:
                 raise ValueError(f'Missing characterization data for pair {(a, b)} in {parameters}')