[XEB] Generalize cirq.google.calibration for XEB (#3881)

Add analogous classes for XEB calibration. This re-configures some of the plumbing to be more general. Most of the structure of the Floquet objects should be untouched / compatible. 

This is a proof of concept that I wanted to get initial comments on from @mrwojtek and @dstrain115 before committing and putting on all the polish (generalize the rest of the functions, tests, notebooks)

Open questions:
 - should the PhasedFSimCalibrationRequest objects have an attribute which is the class of request that they are associated with? 
 - Do you want to go through a deprecation cycle for `prepare_floquet_characterization_for_moment(s)`?

cirq-core/cirq/experiments/__init__.py

@@ -85,3 +85,5 @@
     t2_decay,
     T2DecayResult,
 )
+
+from cirq.experiments.xeb_fitting import XEBPhasedFSimCharacterizationOptions

cirq-core/cirq/experiments/xeb_fitting.py

@@ -37,6 +37,11 @@
     import cirq
     import multiprocessing
 
+    # Workaround for mypy custom dataclasses (python/mypy#5406)
+    from dataclasses import dataclass as json_serializable_dataclass
+else:
+    from cirq.protocols import json_serializable_dataclass
+
 THETA_SYMBOL, ZETA_SYMBOL, CHI_SYMBOL, GAMMA_SYMBOL, PHI_SYMBOL = sympy.symbols(
     'theta zeta chi gamma phi'
 )
@@ -132,7 +137,7 @@ def get_initial_simplex_and_names(
 
 
 # mypy issue: https://github.com/python/mypy/issues/5374
-@dataclass(frozen=True)  # type: ignore
+@json_serializable_dataclass(frozen=True)  # type: ignore
 class XEBPhasedFSimCharacterizationOptions(XEBCharacterizationOptions):
     """Options for calibrating a PhasedFSim-like gate using XEB.
 
@@ -158,11 +163,11 @@ class XEBPhasedFSimCharacterizationOptions(XEBCharacterizationOptions):
     characterize_gamma: bool = True
     characterize_phi: bool = True
 
-    theta_default: float = 0
-    zeta_default: float = 0
-    chi_default: float = 0
-    gamma_default: float = 0
-    phi_default: float = 0
+    theta_default: float = 0.0
+    zeta_default: float = 0.0
+    chi_default: float = 0.0
+    gamma_default: float = 0.0
+    phi_default: float = 0.0
 
     def get_initial_simplex_and_names(
         self, initial_simplex_step_size: float = 0.1
@@ -214,6 +219,14 @@ def get_parameterized_gate(self):
         phi = PHI_SYMBOL if self.characterize_phi else self.phi_default
         return ops.PhasedFSimGate(theta=theta, zeta=zeta, chi=chi, gamma=gamma, phi=phi)
 
+    @staticmethod
+    def should_parameterize(op: 'cirq.Operation') -> bool:
+        if isinstance(op.gate, (ops.PhasedFSimGate, ops.FSimGate)):
+            return True
+        if op.gate == SQRT_ISWAP:
+            return True
+        return False
+
 
 @dataclass(frozen=True)
 class SqrtISwapXEBOptions(XEBPhasedFSimCharacterizationOptions):
@@ -225,10 +238,6 @@ class SqrtISwapXEBOptions(XEBPhasedFSimCharacterizationOptions):
 
     theta_default: float = -np.pi / 4
 
-    @staticmethod
-    def should_parameterize(op: 'cirq.Operation') -> bool:
-        return op.gate == SQRT_ISWAP
-
 
 def parameterize_circuit(
     circuit: 'cirq.Circuit',
@@ -322,7 +331,11 @@ def _mean_infidelity(angles):
     optimization_result = scipy.optimize.minimize(
         _mean_infidelity,
         x0=x0,
-        options={'initial_simplex': initial_simplex, 'xatol': xatol, 'fatol': fatol},
+        options={
+            'initial_simplex': initial_simplex,
+            'xatol': xatol,
+            'fatol': fatol,
+        },
         method='nelder-mead',
     )
 

cirq-core/cirq/experiments/xeb_fitting_test.py

@@ -136,10 +136,11 @@ def _summary_stats(row):
     pd.testing.assert_frame_equal(df_old, df)
 
 
-def test_parameterize_phased_fsim_circuit():
+@pytest.mark.parametrize('gate', [SQRT_ISWAP, cirq.FSimGate(np.pi / 4, 0)])
+def test_parameterize_phased_fsim_circuit(gate):
     q0, q1 = cirq.LineQubit.range(2)
     circuit = rqcg.random_rotations_between_two_qubit_circuit(
-        q0, q1, depth=3, two_qubit_op_factory=lambda a, b, _: SQRT_ISWAP(a, b), seed=52
+        q0, q1, depth=3, two_qubit_op_factory=lambda a, b, _: gate(a, b), seed=52
     )
 
     p_circuit = parameterize_circuit(circuit, SqrtISwapXEBOptions())

cirq-core/cirq/json_resolver_cache.py

@@ -141,6 +141,10 @@ def two_qubit_matrix_gate(matrix):
         'VirtualTag': cirq.VirtualTag,
         'WaitGate': cirq.WaitGate,
         '_QubitAsQid': raw_types._QubitAsQid,
+        # The formatter keeps putting this back
+        # pylint: disable=line-too-long
+        'XEBPhasedFSimCharacterizationOptions': cirq.experiments.XEBPhasedFSimCharacterizationOptions,
+        # pylint: enable=line-too-long
         'XPowGate': cirq.XPowGate,
         'XXPowGate': cirq.XXPowGate,
         'YPowGate': cirq.YPowGate,

cirq-core/cirq/protocols/json_test_data/XEBPhasedFSimCharacterizationOptions.json

@@ -0,0 +1,13 @@
+{
+  "cirq_type": "XEBPhasedFSimCharacterizationOptions",
+  "characterize_theta": true,
+  "characterize_zeta": true,
+  "characterize_chi": true,
+  "characterize_gamma": true,
+  "characterize_phi": true,
+  "theta_default": 0.0,
+  "zeta_default": 0.0,
+  "chi_default": 0.0,
+  "gamma_default": 0.0,
+  "phi_default": 0.0
+}

cirq-core/cirq/protocols/json_test_data/XEBPhasedFSimCharacterizationOptions.repr

@@ -0,0 +1 @@
+cirq.experiments.XEBPhasedFSimCharacterizationOptions(characterize_theta=True, characterize_zeta=True, characterize_chi=True, characterize_gamma=True, characterize_phi=True, theta_default=0.0, zeta_default=0.0, chi_default=0.0, gamma_default=0.0, phi_default=0.0)

cirq-core/requirements.txt

@@ -12,3 +12,4 @@ sortedcontainers~=2.0
 scipy
 sympy
 typing_extensions
+tqdm

cirq-google/cirq_google/__init__.py

@@ -30,14 +30,19 @@
     PhasedFSimCalibrationResult,
     PhasedFSimCharacterization,
     PhasedFSimEngineSimulator,
+    XEBPhasedFSimCalibrationOptions,
+    XEBPhasedFSimCalibrationRequest,
     SQRT_ISWAP_PARAMETERS,
     THETA_ZETA_GAMMA_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     make_zeta_chi_gamma_compensation_for_moments,
     make_zeta_chi_gamma_compensation_for_operations,
     merge_matching_results,
     prepare_floquet_characterization_for_moments,
+    prepare_characterization_for_moments,
     prepare_floquet_characterization_for_moment,
+    prepare_characterization_for_moment,
     prepare_floquet_characterization_for_operations,
+    prepare_characterization_for_operations,
     run_calibrations,
     run_floquet_characterization_for_moments,
     run_zeta_chi_gamma_compensation_for_moments,

cirq-google/cirq_google/calibration/__init__.py

@@ -24,6 +24,8 @@
     PhasedFSimCalibrationRequest,
     PhasedFSimCalibrationResult,
     PhasedFSimCharacterization,
+    XEBPhasedFSimCalibrationOptions,
+    XEBPhasedFSimCalibrationRequest,
     SQRT_ISWAP_PARAMETERS,
     THETA_ZETA_GAMMA_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
@@ -36,8 +38,11 @@
     make_zeta_chi_gamma_compensation_for_moments,
     make_zeta_chi_gamma_compensation_for_operations,
     prepare_floquet_characterization_for_moments,
+    prepare_characterization_for_moments,
     prepare_floquet_characterization_for_moment,
+    prepare_characterization_for_moment,
     prepare_floquet_characterization_for_operations,
+    prepare_characterization_for_operations,
     run_calibrations,
     run_floquet_characterization_for_moments,
     run_zeta_chi_gamma_compensation_for_moments,

cirq-google/cirq_google/calibration/engine_simulator_test.py

@@ -36,6 +36,7 @@ def test_test_calibration_request():
     request = TestPhasedFSimCalibrationRequest(
         gate=cirq.FSimGate(np.pi / 4, 0.5),
         pairs=((a, b),),
+        options=ALL_ANGLES_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     )
 
     assert request.to_calibration_layer() is NotImplemented
@@ -106,6 +107,7 @@ def test_floquet_get_calibrations_when_invalid_request_fails() -> None:
                 TestPhasedFSimCalibrationRequest(
                     gate=cirq.FSimGate(np.pi / 4, 0.5),
                     pairs=((a, b),),
+                    options=ALL_ANGLES_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
                 )
             ]
         )