Support creating XEB calibration options for arbitary FSim like gates (…

…quantumlib#6657)

This PR supports creating FSim calibration options for arbitary FSim like gates. It also split the XEB function in two in prerpartion for implmenting XEB calibration for arbitary gates.

part of breaking quantumlib#6568 into smaller PRs.

cirq/experiments/two_qubit_xeb.py

@@ -347,7 +347,7 @@ def plot_histogram(
         return ax
 
 
-def parallel_two_qubit_xeb(
+def parallel_xeb_workflow(
     sampler: 'cirq.Sampler',
     qubits: Optional[Sequence['cirq.GridQubit']] = None,
     entangling_gate: 'cirq.Gate' = ops.CZ,
@@ -358,8 +358,8 @@ def parallel_two_qubit_xeb(
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     ax: Optional[plt.Axes] = None,
     **plot_kwargs,
-) -> TwoQubitXEBResult:
-    """A convenience method that runs the full XEB workflow.
+) -> Tuple[pd.DataFrame, Sequence['cirq.Circuit'], pd.DataFrame]:
+    """A utility method that runs the full XEB workflow.
 
     Args:
         sampler: The quantum engine or simulator to run the circuits.
@@ -375,7 +375,12 @@ def parallel_two_qubit_xeb(
         **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
 
     Returns:
-        A TwoQubitXEBResult object representing the results of the experiment.
+        - A DataFrame with columns 'cycle_depth' and 'fidelity'.
+        - The circuits used to perform XEB.
+        - A pandas dataframe with index given by ['circuit_i', 'cycle_depth'].
+            Columns always include "sampled_probs". If `combinations_by_layer` is
+            not `None` and you are doing parallel XEB, additional metadata columns
+            will be attached to the returned DataFrame.
 
     Raises:
         ValueError: If qubits are not specified and the sampler has no device.
@@ -420,6 +425,52 @@ def parallel_two_qubit_xeb(
         sampled_df=sampled_df, circuits=circuit_library, cycle_depths=cycle_depths
     )
 
+    return fids, circuit_library, sampled_df
+
+
+def parallel_two_qubit_xeb(
+    sampler: 'cirq.Sampler',
+    qubits: Optional[Sequence['cirq.GridQubit']] = None,
+    entangling_gate: 'cirq.Gate' = ops.CZ,
+    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,
+    ax: Optional[plt.Axes] = None,
+    **plot_kwargs,
+) -> TwoQubitXEBResult:
+    """A convenience method that runs the full XEB workflow.
+
+    Args:
+        sampler: The quantum engine or simulator to run the circuits.
+        qubits: Qubits under test. If none, uses all qubits on the sampler's device.
+        entangling_gate: The entangling gate to use.
+        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.
+        ax: the plt.Axes to plot the device layout on. If not given,
+            no plot is created.
+        **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
+    Returns:
+        A TwoQubitXEBResult object representing the results of the experiment.
+    Raises:
+        ValueError: If qubits are not specified and the sampler has no device.
+    """
+    fids, *_ = parallel_xeb_workflow(
+        sampler=sampler,
+        qubits=qubits,
+        entangling_gate=entangling_gate,
+        n_repetitions=n_repetitions,
+        n_combinations=n_combinations,
+        n_circuits=n_circuits,
+        cycle_depths=cycle_depths,
+        random_state=random_state,
+        ax=ax,
+        **plot_kwargs,
+    )
     return TwoQubitXEBResult(fit_exponential_decays(fids))
 
 

cirq/experiments/xeb_fitting.py

@@ -146,6 +146,69 @@ def get_initial_simplex_and_names(
         """Return an initial Nelder-Mead simplex and the names for each parameter."""
 
 
+def _try_defaults_from_unitary(gate: 'cirq.Gate') -> Optional[Dict[str, 'cirq.TParamVal']]:
+    r"""Try to figure out the PhasedFSim angles from the unitary of the gate.
+
+    The unitary of a PhasedFSimGate has the form:
+    $$
+    \begin{bmatrix}
+        1 & 0 & 0 & 0 \\
+        0 & e^{-i \gamma - i \zeta} \cos(\theta) & -i e^{-i \gamma + i\chi} \sin(\theta) & 0 \\
+        0 & -i e^{-i \gamma - i \chi} \sin(\theta) & e^{-i \gamma + i \zeta} \cos(\theta) & 0 \\
+        0 & 0 & 0 & e^{-2i \gamma - i \phi}
+    \end{bmatrix}
+    $$
+    That's the information about the five angles $\theta, \phi, \gamma, \zeta, \chi$ is encoded in
+    the submatrix unitary[1:3, 1:3] and the element u[3][3]. With some algebra, we can isolate each
+    of the angles as an argument of a combination of those elements (and potentially other angles).
+
+    Args:
+        A cirq gate.
+
+    Returns:
+        A dictionary mapping angles to values or None if the gate doesn't have a unitary or if it
+        can't be represented by a PhasedFSimGate.
+    """
+    u = protocols.unitary(gate, default=None)
+    if u is None:
+        return None
+
+    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])
+
+    if np.allclose(
+        u,
+        protocols.unitary(
+            ops.PhasedFSimGate(theta=theta, phi=phi, chi=chi, zeta=zeta, gamma=gamma)
+        ),
+    ):
+        return {
+            'theta_default': theta,
+            'phi_default': phi,
+            'gamma_default': gamma,
+            'zeta_default': zeta,
+            'chi_default': chi,
+        }
+    return None
+
+
 def phased_fsim_angles_from_gate(gate: 'cirq.Gate') -> Dict[str, 'cirq.TParamVal']:
     """For a given gate, return a dictionary mapping '{angle}_default' to its noiseless value
     for the five PhasedFSim angles."""
@@ -175,6 +238,11 @@ def phased_fsim_angles_from_gate(gate: 'cirq.Gate') -> Dict[str, 'cirq.TParamVal
             'phi_default': gate.phi,
         }
 
+    # Handle all gates that can be represented using an FSimGate.
+    from_unitary = _try_defaults_from_unitary(gate)
+    if from_unitary is not None:
+        return from_unitary
+
     raise ValueError(f"Unknown default angles for {gate}.")
 
 
@@ -580,15 +648,6 @@ def _fit_exponential_decay(
     return a, layer_fid, a_std, layer_fid_std
 
 
-def _one_unique(df, name, default):
-    """Helper function to assert that there's one unique value in a column and return it."""
-    if name not in df.columns:
-        return default
-    vals = df[name].unique()
-    assert len(vals) == 1, name
-    return vals[0]
-
-
 def fit_exponential_decays(fidelities_df: pd.DataFrame) -> pd.DataFrame:
     """Fit exponential decay curves to a fidelities DataFrame.
 

cirq/experiments/xeb_fitting_test.py

@@ -32,6 +32,7 @@
     fit_exponential_decays,
     before_and_after_characterization,
     XEBPhasedFSimCharacterizationOptions,
+    phased_fsim_angles_from_gate,
 )
 from cirq.experiments.xeb_sampling import sample_2q_xeb_circuits
 
@@ -354,7 +355,7 @@ def test_options_with_defaults_from_gate():
     assert options.zeta_default == 0.0
 
     with pytest.raises(ValueError):
-        _ = XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(cirq.CZ)
+        _ = XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(cirq.XX)
 
 
 def test_options_defaults_set():
@@ -395,3 +396,34 @@ def test_options_defaults_set():
         phi_default=0.0,
     )
     assert o3.defaults_set() is True
+
+
+def _random_angles(n, seed):
+    rng = np.random.default_rng(seed)
+    r = 2 * rng.random((n, 5)) - 1
+    return np.pi * r
+
+
+@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,
+    ]
+    + [cirq.PhasedFSimGate(*r) for r in _random_angles(10, 0)],
+)
+def test_phased_fsim_angles_from_gate(gate):
+    angles = phased_fsim_angles_from_gate(gate)
+    angles = {k.removesuffix('_default'): v for k, v in angles.items()}
+    phasedfsim = cirq.PhasedFSimGate(**angles)
+    np.testing.assert_allclose(cirq.unitary(phasedfsim), cirq.unitary(gate), atol=1e-9)
+
+
+def test_phased_fsim_angles_from_gate_unsupporet_gate():
+    with pytest.raises(ValueError, match='Unknown default angles'):
+        _ = phased_fsim_angles_from_gate(cirq.testing.TwoQubitGate())