Create a workflow to perform FSim based z-phase calibration

cirq-core/cirq/experiments/__init__.py

@@ -71,3 +71,6 @@
     parallel_two_qubit_xeb,
     run_rb_and_xeb,
 )
+
+
+from cirq.experiments.z_phase_calibration import calibrate_z_phases

cirq-core/cirq/experiments/two_qubit_xeb.py

@@ -17,7 +17,6 @@
 
 from dataclasses import dataclass
 from types import MappingProxyType
-import itertools
 import functools
 
 from matplotlib import pyplot as plt
@@ -27,13 +26,13 @@
 
 from cirq import ops, value, vis
 from cirq.experiments.xeb_sampling import sample_2q_xeb_circuits
-from cirq.experiments.xeb_fitting import benchmark_2q_xeb_fidelities
-from cirq.experiments.xeb_fitting import fit_exponential_decays, exponential_decay
+from cirq.experiments import xeb_fitting as xeb_fitting
 from cirq.experiments import random_quantum_circuit_generation as rqcg
 from cirq.experiments.qubit_characterizations import (
     ParallelRandomizedBenchmarkingResult,
     parallel_single_qubit_randomized_benchmarking,
 )
+from cirq.experiments.xeb_utils import grid_qubits_to_graph
 from cirq.qis import noise_utils
 from cirq._compat import cached_method
 
@@ -48,10 +47,6 @@ def _grid_qubits_for_sampler(sampler: 'cirq.Sampler') -> Optional[Sequence['cirq
     return None
 
 
-def _manhattan_distance(qubit1: 'cirq.GridQubit', qubit2: 'cirq.GridQubit') -> int:
-    return abs(qubit1.row - qubit2.row) + abs(qubit1.col - qubit2.col)
-
-
 @dataclass(frozen=True)
 class TwoQubitXEBResult:
     """Results from an XEB experiment."""
@@ -124,7 +119,7 @@ def plot_fitted_exponential(
         depths = np.linspace(0, np.max(record['cycle_depths']))
         ax.plot(
             depths,
-            exponential_decay(depths, a=record['a'], layer_fid=record['layer_fid']),
+            xeb_fitting.exponential_decay(depths, a=record['a'], layer_fid=record['layer_fid']),
             label='estimated exponential decay',
             **plot_kwargs,
         )
@@ -347,7 +342,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 +353,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 +370,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.
@@ -387,9 +387,7 @@ def parallel_two_qubit_xeb(
         if qubits is None:
             raise ValueError("Couldn't determine qubits from sampler. Please specify them.")
 
-    graph = nx.Graph(
-        pair for pair in itertools.combinations(qubits, 2) if _manhattan_distance(*pair) == 1
-    )
+    graph = grid_qubits_to_graph(qubits)
 
     if ax is not None:
         nx.draw_networkx(graph, pos={q: (q.row, q.col) for q in qubits}, ax=ax)
@@ -416,11 +414,59 @@ def parallel_two_qubit_xeb(
         repetitions=n_repetitions,
     )
 
-    fids = benchmark_2q_xeb_fidelities(
+    fids = xeb_fitting.benchmark_2q_xeb_fidelities(
         sampled_df=sampled_df, circuits=circuit_library, cycle_depths=cycle_depths
     )
 
-    return TwoQubitXEBResult(fit_exponential_decays(fids))
+    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(xeb_fitting.fit_exponential_decays(fids))
 
 
 def run_rb_and_xeb(

cirq-core/cirq/experiments/xeb_fitting.py

@@ -14,17 +14,20 @@
 """Estimation of fidelity associated with experimental circuit executions."""
 import dataclasses
 from abc import abstractmethod, ABC
-from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING
+from typing import Dict, Iterable, List, Union, Optional, Sequence, Tuple, TYPE_CHECKING
 
+import tqdm
 import numpy as np
 import pandas as pd
 import sympy
 from cirq import circuits, ops, protocols, _import
 from cirq.experiments.xeb_simulation import simulate_2q_xeb_circuits
+from cirq.experiments import xeb_utils
 
 if TYPE_CHECKING:
     import cirq
     import multiprocessing
+    import concurrent.futures
     import scipy.optimize
 
 # We initialize these lazily, otherwise they slow global import speed.
@@ -43,12 +46,26 @@ def benchmark_2q_xeb_fidelities(
     param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
     pool: Optional['multiprocessing.pool.Pool'] = None,
 ) -> pd.DataFrame:
-    """Simulate and benchmark two-qubit XEB circuits.
+    r"""Simulate and benchmark two-qubit XEB circuits.
 
     This uses the estimator from
     `cirq.experiments.fidelity_estimation.least_squares_xeb_fidelity_from_expectations`, but
     adapted for use on pandas DataFrames for efficient vectorized operation.
 
+    The error model is assumed to be a channel that applies the unitary with probability $f$
+    (i.e. fidelity) or does nothing. This mapping is represented by
+    $$
+        \ket{\psi} \xrightarrow \rho_U = f \ket{\psi_U}\bra{\psi_U} + (1 - f) I / D
+    $$
+    Where $\rho_U$ is the density matrix after the operation. This leads to
+    $$
+        Tr(\rho_U O_U) = f \bra{\psi_U} O_U \ket{\psi_U} + (1 - f) Tr(O_U / D)
+    $$
+    setting $m_U = Tr(\rho_U O_U), u_U = Tr(O_U / D), e_U = \bra{\psi_U} O_U \ket{\psi_U}$ we get
+    $$
+        f = \frac{m_U - u_U}{e_U - u_U}
+    $$
+
     Args:
         sampled_df: The sampled results to benchmark. This is likely produced by a call to
             `sample_2q_xeb_circuits`.
@@ -92,6 +109,8 @@ def benchmark_2q_xeb_fidelities(
     D = 4  # two qubits
     pure_probs = np.array(df['pure_probs'].to_list())
     sampled_probs = np.array(df['sampled_probs'].to_list())
+    # pure_probs = np.sqrt(np.array(df['pure_probs'].to_list()))
+    # sampled_probs = np.sqrt(np.array(df['sampled_probs'].to_list()))
     df['e_u'] = np.sum(pure_probs**2, axis=1)
     df['u_u'] = np.sum(pure_probs, axis=1) / D
     df['m_u'] = np.sum(pure_probs * sampled_probs, axis=1)
@@ -130,11 +149,6 @@ def _try_keep(k):
 
 
 class XEBCharacterizationOptions(ABC):
-    @staticmethod
-    @abstractmethod
-    def should_parameterize(op: 'cirq.Operation') -> bool:
-        """Whether to replace `op` with a parameterized version."""
-
     @abstractmethod
     def get_parameterized_gate(self) -> 'cirq.Gate':
         """The parameterized gate to use."""
@@ -175,6 +189,15 @@ def phased_fsim_angles_from_gate(gate: 'cirq.Gate') -> Dict[str, 'cirq.TParamVal
             'phi_default': gate.phi,
         }
 
+    # Handle all gates that preserve excitations (i.e. fermionic gates).
+    u = protocols.unitary(gate)
+    phi = -np.angle(u[3, 3])
+    theta = np.angle(u[1, 1] - u[1, 2])
+    if np.allclose(u, protocols.unitary(ops.FSimGate(theta=theta, phi=phi))):
+        defaults['theta_default'] = theta
+        defaults['phi_default'] = phi
+        return defaults
+
     raise ValueError(f"Unknown default angles for {gate}.")
 
 
@@ -266,10 +289,6 @@ 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:
-        return isinstance(op.gate, (ops.PhasedFSimGate, ops.ISwapPowGate, ops.FSimGate))
-
     def defaults_set(self) -> bool:
         """Whether the default angles are set.
 
@@ -317,17 +336,18 @@ def SqrtISwapXEBOptions(*args, **kwargs):
 
 
 def parameterize_circuit(
-    circuit: 'cirq.Circuit', options: XEBCharacterizationOptions
+    circuit: 'cirq.Circuit',
+    options: XEBCharacterizationOptions,
+    target: Union[ops.GateFamily, ops.Gateset] = ops.Gateset(
+        ops.PhasedFSimGate, ops.ISwapPowGate, ops.FSimGate
+    ),
 ) -> 'cirq.Circuit':
     """Parameterize PhasedFSim-like gates in a given circuit according to
     `phased_fsim_options`.
     """
     gate = options.get_parameterized_gate()
     return circuits.Circuit(
-        circuits.Moment(
-            gate.on(*op.qubits) if options.should_parameterize(op) else op
-            for op in moment.operations
-        )
+        circuits.Moment(gate.on(*op.qubits) if op in target else op for op in moment.operations)
         for moment in circuit.moments
     )
 
@@ -398,7 +418,6 @@ def _mean_infidelity(angles):
         fids = benchmark_2q_xeb_fidelities(
             sampled_df, parameterized_circuits, cycle_depths, param_resolver=params, pool=pool
         )
-
         loss = 1 - fids['fidelity'].mean()
         if verbose:
             print(f"Loss: {loss:7.3g}", flush=True)
@@ -456,7 +475,9 @@ def characterize_phased_fsim_parameters_with_xeb_by_pair(
     initial_simplex_step_size: float = 0.1,
     xatol: float = 1e-3,
     fatol: float = 1e-3,
-    pool: Optional['multiprocessing.pool.Pool'] = None,
+    pool: Optional[
+        Union['multiprocessing.pool.Pool', 'concurrent.futures.ThreadPoolExecutor']
+    ] = None,
 ) -> XEBCharacterizationResult:
     """Run a classical optimization to fit phased fsim parameters to experimental data, and
     thereby characterize PhasedFSim-like gates grouped by pairs.
@@ -493,11 +514,13 @@ def characterize_phased_fsim_parameters_with_xeb_by_pair(
         fatol=fatol,
     )
     subselected_dfs = [sampled_df[sampled_df['pair'] == pair] for pair in pairs]
-    if pool is not None:
-        results = pool.map(closure, subselected_dfs)
-    else:
-        results = [closure(df) for df in subselected_dfs]
-
+    results = xeb_utils.execute_with_progress_par(
+        closure,
+        subselected_dfs,
+        pool=pool,
+        progress_bar=tqdm.tqdm,
+        # desc='characterize fsim parameters',
+    )
     optimization_results = {}
     all_final_params = {}
     fid_dfs = []
@@ -513,7 +536,7 @@ def characterize_phased_fsim_parameters_with_xeb_by_pair(
     )
 
 
-def exponential_decay(cycle_depths: np.ndarray, a: float, layer_fid: float) -> np.ndarray:
+def exponential_decay(cycle_depths: np.ndarray, a: float, layer_fid: float, b: float) -> np.ndarray:
     """An exponential decay for fitting.
 
     This computes `a * layer_fid**cycle_depths`
@@ -522,12 +545,13 @@ def exponential_decay(cycle_depths: np.ndarray, a: float, layer_fid: float) -> n
         cycle_depths: The various depths at which fidelity was estimated. This is the independent
             variable in the exponential function.
         a: A scale parameter in the exponential function.
+        b: An offset parameter.
         layer_fid: The base of the exponent in the exponential function.
     """
-    return a * layer_fid**cycle_depths
+    return a * layer_fid**cycle_depths + b
 
 
-def _fit_exponential_decay(
+def fit_exponential_decay(
     cycle_depths: np.ndarray, fidelities: np.ndarray
 ) -> Tuple[float, float, float, float]:
     """Fit an exponential model fidelity = a * layer_fid**x using nonlinear least squares.
@@ -566,29 +590,20 @@ def _fit_exponential_decay(
     a_0 = np.clip(np.exp(intercept), 0, 1)
 
     try:
-        (a, layer_fid), pcov = optimize.curve_fit(
+        (a, layer_fid, _), pcov = optimize.curve_fit(
             exponential_decay,
             cycle_depths,
             fidelities,
-            p0=(a_0, layer_fid_0),
-            bounds=((0, 0), (1, 1)),
+            p0=(a_0, layer_fid_0, 0.999),
+            bounds=((0, 0, 0), (1, 1, 1)),
         )
     except ValueError:  # pragma: no cover
         return 0, 0, np.inf, np.inf
 
-    a_std, layer_fid_std = np.sqrt(np.diag(pcov))
+    a_std, layer_fid_std = np.sqrt(np.diag(pcov[:2]))
     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.
 
@@ -605,7 +620,7 @@ def fit_exponential_decays(fidelities_df: pd.DataFrame) -> pd.DataFrame:
     """
 
     def _per_pair(f1):
-        a, layer_fid, a_std, layer_fid_std = _fit_exponential_decay(
+        a, layer_fid, a_std, layer_fid_std = fit_exponential_decay(
             f1['cycle_depth'], f1['fidelity']
         )
         record = {