[XEB] Allow default angles to be inferred from a gate (#4092)

Instead of defaulting angles to 0.0 -- which is a very sharp edge that can cause optimization to fail -- default to `None` and either require the user fully specify the angle settings or provide a gate to get sensible defaults from. In `xeb_wrapper` this gate is already around as the calibration request gate.

cirq-core/cirq/experiments/xeb_fitting.py

@@ -21,6 +21,7 @@
     Tuple,
     TYPE_CHECKING,
     Dict,
+    Iterable,
 )
 
 import numpy as np
@@ -136,6 +137,38 @@ def get_initial_simplex_and_names(
         """Return an initial Nelder-Mead simplex and the names for each parameter."""
 
 
+def phased_fsim_angles_from_gate(gate: 'cirq.Gate') -> Dict[str, float]:
+    """For a given gate, return a dictionary mapping '{angle}_default' to its noiseless value
+    for the five PhasedFSim angles."""
+    defaults = {
+        'theta_default': 0.0,
+        'zeta_default': 0.0,
+        'chi_default': 0.0,
+        'gamma_default': 0.0,
+        'phi_default': 0.0,
+    }
+    if gate == SQRT_ISWAP:
+        defaults['theta_default'] = -np.pi / 4
+        return defaults
+    if gate == SQRT_ISWAP ** -1:
+        defaults['theta_default'] = np.pi / 4
+        return defaults
+    if isinstance(gate, ops.FSimGate):
+        defaults['theta_default'] = gate.theta
+        defaults['phi_default'] = gate.phi
+        return defaults
+    if isinstance(gate, ops.PhasedFSimGate):
+        return {
+            'theta_default': gate.theta,
+            'zeta_default': gate.zeta,
+            'chi_default': gate.chi,
+            'gamma_default': gate.gamma,
+            'phi_default': gate.phi,
+        }
+
+    raise ValueError(f"Unknown default angles for {gate}.")
+
+
 # mypy issue: https://github.com/python/mypy/issues/5374
 @json_serializable_dataclass(frozen=True)  # type: ignore
 class XEBPhasedFSimCharacterizationOptions(XEBCharacterizationOptions):
@@ -163,11 +196,27 @@ class XEBPhasedFSimCharacterizationOptions(XEBCharacterizationOptions):
     characterize_gamma: bool = True
     characterize_phi: bool = True
 
-    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
+    theta_default: Optional[float] = None
+    zeta_default: Optional[float] = None
+    chi_default: Optional[float] = None
+    gamma_default: Optional[float] = None
+    phi_default: Optional[float] = None
+
+    def _iter_angles(self) -> Iterable[Tuple[bool, Optional[float], 'sympy.Symbol']]:
+        yield from (
+            (self.characterize_theta, self.theta_default, THETA_SYMBOL),
+            (self.characterize_zeta, self.zeta_default, ZETA_SYMBOL),
+            (self.characterize_chi, self.chi_default, CHI_SYMBOL),
+            (self.characterize_gamma, self.gamma_default, GAMMA_SYMBOL),
+            (self.characterize_phi, self.phi_default, PHI_SYMBOL),
+        )
+
+    def _iter_angles_for_characterization(self) -> Iterable[Tuple[Optional[float], 'sympy.Symbol']]:
+        yield from (
+            (default, symbol)
+            for characterize, default, symbol in self._iter_angles()
+            if characterize
+        )
 
     def get_initial_simplex_and_names(
         self, initial_simplex_step_size: float = 0.1
@@ -185,21 +234,12 @@ def get_initial_simplex_and_names(
         """
         x0 = []
         names = []
-        if self.characterize_theta:
-            x0 += [self.theta_default]
-            names += [THETA_SYMBOL.name]
-        if self.characterize_zeta:
-            x0 += [self.zeta_default]
-            names += [ZETA_SYMBOL.name]
-        if self.characterize_chi:
-            x0 += [self.chi_default]
-            names += [CHI_SYMBOL.name]
-        if self.characterize_gamma:
-            x0 += [self.gamma_default]
-            names += [GAMMA_SYMBOL.name]
-        if self.characterize_phi:
-            x0 += [self.phi_default]
-            names += [PHI_SYMBOL.name]
+
+        for default, symbol in self._iter_angles_for_characterization():
+            if default is None:
+                raise ValueError(f'{symbol.name}_default was not set.')
+            x0.append(default)
+            names.append(symbol.name)
 
         x0 = np.asarray(x0)
         n_param = len(x0)
@@ -223,16 +263,45 @@ def get_parameterized_gate(self):
     def should_parameterize(op: 'cirq.Operation') -> bool:
         return isinstance(op.gate, (ops.PhasedFSimGate, ops.ISwapPowGate, ops.FSimGate))
 
+    def defaults_set(self) -> bool:
+        """Whether the default angles are set.
 
-@dataclass(frozen=True)
-class SqrtISwapXEBOptions(XEBPhasedFSimCharacterizationOptions):
-    """Options for calibrating a sqrt(ISWAP) gate using XEB.
+        This only considers angles where characterize_{angle} is True. If all such angles have
+        {angle}_default set to a value, this returns True. If none of the defaults are set,
+        this returns False. If some defaults are set, we raise an exception.
+        """
+        defaults_set = [default is not None for _, default, _ in self._iter_angles()]
+        if any(defaults_set):
+            if all(defaults_set):
+                return True
+
+            problems = [
+                symbol.name for _, default, symbol in self._iter_angles() if default is None
+            ]
+            raise ValueError(f"Some angles are set, but values for {problems} are not.")
+        return False
+
+    def with_defaults_from_gate(
+        self, gate: 'cirq.Gate', gate_to_angles_func=phased_fsim_angles_from_gate
+    ):
+        """A new Options class with {angle}_defaults inferred from `gate`.
+
+        This keeps the same settings for the characterize_{angle} booleans, but will disregard
+        any current {angle}_default values.
+        """
+        return XEBPhasedFSimCharacterizationOptions(
+            characterize_theta=self.characterize_theta,
+            characterize_zeta=self.characterize_zeta,
+            characterize_chi=self.characterize_chi,
+            characterize_gamma=self.characterize_gamma,
+            characterize_phi=self.characterize_phi,
+            **gate_to_angles_func(gate),
+        )
 
-    As such, the default for theta is changed to -pi/4 and the parameterization
-    predicate seeks out sqrt(ISWAP) gates.
-    """
 
-    theta_default: float = -np.pi / 4
+def SqrtISwapXEBOptions(*args, **kwargs):
+    """Options for calibrating a sqrt(ISWAP) gate using XEB."""
+    return XEBPhasedFSimCharacterizationOptions(*args, **kwargs).with_defaults_from_gate(SQRT_ISWAP)
 
 
 def parameterize_circuit(

cirq-core/cirq/experiments/xeb_fitting_test.py

@@ -32,6 +32,7 @@
     _fit_exponential_decay,
     fit_exponential_decays,
     before_and_after_characterization,
+    XEBPhasedFSimCharacterizationOptions,
 )
 from cirq.experiments.xeb_sampling import sample_2q_xeb_circuits
 
@@ -315,3 +316,66 @@ def test_fit_exponential_decays_negative_fids():
     assert layer_fid == 0
     assert a_std == np.inf
     assert layer_fid_std == np.inf
+
+
+def test_options_with_defaults_from_gate():
+    options = XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(cirq.ISWAP ** 0.5)
+    np.testing.assert_allclose(options.theta_default, -np.pi / 4)
+    options = XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(cirq.ISWAP ** -0.5)
+    np.testing.assert_allclose(options.theta_default, np.pi / 4)
+
+    options = XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(
+        cirq.FSimGate(0.1, 0.2)
+    )
+    assert options.theta_default == 0.1
+    assert options.phi_default == 0.2
+
+    options = XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(
+        cirq.PhasedFSimGate(0.1)
+    )
+    assert options.theta_default == 0.1
+    assert options.phi_default == 0.0
+    assert options.zeta_default == 0.0
+
+    with pytest.raises(ValueError):
+        _ = XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(cirq.CZ)
+
+
+def test_options_defaults_set():
+    o1 = XEBPhasedFSimCharacterizationOptions(
+        characterize_zeta=True,
+        characterize_chi=True,
+        characterize_gamma=True,
+        characterize_theta=False,
+        characterize_phi=False,
+    )
+    assert o1.defaults_set() is False
+    with pytest.raises(ValueError):
+        o1.get_initial_simplex_and_names()
+
+    o2 = XEBPhasedFSimCharacterizationOptions(
+        characterize_zeta=True,
+        characterize_chi=True,
+        characterize_gamma=True,
+        characterize_theta=False,
+        characterize_phi=False,
+        zeta_default=0.1,
+        chi_default=0.2,
+        gamma_default=0.3,
+    )
+    with pytest.raises(ValueError):
+        _ = o2.defaults_set()
+
+    o3 = XEBPhasedFSimCharacterizationOptions(
+        characterize_zeta=True,
+        characterize_chi=True,
+        characterize_gamma=True,
+        characterize_theta=False,
+        characterize_phi=False,
+        zeta_default=0.1,
+        chi_default=0.2,
+        gamma_default=0.3,
+        theta_default=0.0,
+        phi_default=0.0,
+    )
+    assert o3.defaults_set() is True

cirq-google/cirq_google/calibration/phased_fsim.py

@@ -485,7 +485,9 @@ def to_args(self) -> Dict[str, Any]:
         if self.xatol is not None:
             args['xatol'] = self.xatol
 
-        args.update(dataclasses.asdict(self.fsim_options))
+        fsim_options = dataclasses.asdict(self.fsim_options)
+        fsim_options = {k: v for k, v in fsim_options.items() if v is not None}
+        args.update(fsim_options)
         return args
 
     def create_phased_fsim_request(

cirq-google/cirq_google/calibration/phased_fsim_test.py

@@ -183,11 +183,6 @@ def test_xeb_to_calibration_layer():
             'characterize_chi': False,
             'characterize_gamma': False,
             'characterize_phi': True,
-            'theta_default': 0.0,
-            'zeta_default': 0.0,
-            'chi_default': 0.0,
-            'gamma_default': 0.0,
-            'phi_default': 0.0,
         },
     )
 

cirq-google/cirq_google/calibration/test_data/xeb_calibration_layer.textproto

@@ -71,14 +71,6 @@ args {
     }
   }
 }
-args {
-  key: "chi_default"
-  value {
-    arg_value {
-      float_value: 0.0
-    }
-  }
-}
 args {
   key: "cycle_depths"
   value {
@@ -95,14 +87,6 @@ args {
     }
   }
 }
-args {
-  key: "gamma_default"
-  value {
-    arg_value {
-      float_value: 0.0
-    }
-  }
-}
 args {
   key: "n_combinations"
   value {
@@ -119,22 +103,6 @@ args {
     }
   }
 }
-args {
-  key: "phi_default"
-  value {
-    arg_value {
-      float_value: 0.0
-    }
-  }
-}
-args {
-  key: "theta_default"
-  value {
-    arg_value {
-      float_value: 0.0
-    }
-  }
-}
 args {
   key: "xatol"
   value {
@@ -143,11 +111,3 @@ args {
     }
   }
 }
-args {
-  key: "zeta_default"
-  value {
-    arg_value {
-      float_value: 0.0
-    }
-  }
-}

cirq-google/cirq_google/calibration/xeb_wrapper.py

@@ -86,15 +86,20 @@ def run_local_xeb_calibration(
     # )
 
     # 5. Characterize by fitting angles.
-    pcircuits = [xebf.parameterize_circuit(circuit, options.fsim_options) for circuit in circuits]
+    if options.fsim_options.defaults_set():
+        fsim_options = options.fsim_options
+    else:
+        fsim_options = options.fsim_options.with_defaults_from_gate(calibration.gate)
+
+    pcircuits = [xebf.parameterize_circuit(circuit, fsim_options) for circuit in circuits]
     fatol = options.fatol if options.fatol is not None else 5e-3
     xatol = options.xatol if options.xatol is not None else 5e-3
     with _maybe_multiprocessing_pool(n_processes=options.n_processes) as pool:
         char_results = xebf.characterize_phased_fsim_parameters_with_xeb_by_pair(
             sampled_df=sampled_df,
             parameterized_circuits=pcircuits,
             cycle_depths=cycle_depths,
-            options=options.fsim_options,
+            options=fsim_options,
             pool=pool,
             fatol=fatol,
             xatol=xatol,