[XEB] Local wrapper for characterization functionality (quantumlib#4047)

XEB calibration can be run through the calibration api, i.e. by a call to `engine.run_calibration` but since all the code is open-sourced in Cirq and only requires the ability to sample from a quantum processor, we can write a wrapper which can be dispatched to by `cg.workflow.run_calibrations` that submits circuits using a `cirq.Sampler` and performs XEB angle optimization locally.

cirq-core/cirq/experiments/xeb_fitting.py

@@ -54,7 +54,7 @@ def benchmark_2q_xeb_fidelities(
     cycle_depths: Sequence[int],
     param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
     pool: Optional['multiprocessing.pool.Pool'] = None,
-):
+) -> pd.DataFrame:
     """Simulate and benchmark two-qubit XEB circuits.
 
     This uses the estimator from
@@ -221,11 +221,7 @@ def get_parameterized_gate(self):
 
     @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
+        return isinstance(op.gate, (ops.PhasedFSimGate, ops.ISwapPowGate, ops.FSimGate))
 
 
 @dataclass(frozen=True)

cirq-google/cirq_google/calibration/phased_fsim.py

@@ -247,6 +247,7 @@ class PhasedFSimCalibrationResult:
             quibts a and b either only (a, b) or only (b, a) is present.
         gate: Characterized gate for each qubit pair. This is copied from the matching
             PhasedFSimCalibrationRequest and is included to preserve execution context.
+        options: The options used to gather this result.
     """
 
     parameters: Dict[Tuple[Qid, Qid], PhasedFSimCharacterization]
@@ -467,6 +468,52 @@ def _from_json_dict_(cls, **kwargs):
         return cls(**kwargs)
 
 
+@json_serializable_dataclass(frozen=True)
+class LocalXEBPhasedFSimCalibrationOptions(XEBPhasedFSimCalibrationOptions):
+    """Options for configuring a PhasedFSim calibration using a local version of XEB.
+
+    XEB uses the fidelity of random circuits to characterize PhasedFSim gates. The parameters
+    of the gate are varied by a classical optimizer to maximize the observed fidelities.
+
+    These "Local" options (corresponding to `LocalXEBPhasedFSimCalibrationRequest`) instruct
+    `cirq_google.run_calibrations` to execute XEB analysis locally (not via the quantum
+    engine). As such, `run_calibrations` can work with any `cirq.Sampler`, not just
+    `QuantumEngineSampler`.
+
+    Args:
+        n_library_circuits: The number of distinct, two-qubit random circuits to use in our
+            library of random circuits. This should be the same order of magnitude as
+            `n_combinations`.
+        n_combinations: We take each library circuit and randomly assign it to qubit pairs.
+            This parameter controls the number of random combinations of the two-qubit random
+            circuits we execute. Higher values increase the precision of estimates but linearly
+            increase experimental runtime.
+        cycle_depths: We run the random circuits at these cycle depths to fit an exponential
+            decay in the fidelity.
+        fatol: The absolute convergence tolerance for the objective function evaluation in
+            the Nelder-Mead optimization. This controls the runtime of the classical
+            characterization optimization loop.
+        xatol: The absolute convergence tolerance for the parameter estimates in
+            the Nelder-Mead optimization. This controls the runtime of the classical
+            characterization optimization loop.
+        fsim_options: An instance of `XEBPhasedFSimCharacterizationOptions` that controls aspects
+            of the PhasedFSim characterization like initial guesses and which angles to
+            characterize.
+        n_processes: The number of multiprocessing processes to analyze the XEB characterization
+            data. By default, we use a value equal to the number of CPU cores. If `1` is specified,
+            multiprocessing is not used.
+    """
+
+    n_processes: Optional[int] = None
+
+    def create_phased_fsim_request(
+        self,
+        pairs: Tuple[Tuple[Qid, Qid], ...],
+        gate: Gate,
+    ):
+        return LocalXEBPhasedFSimCalibrationRequest(pairs=pairs, gate=gate, options=self)
+
+
 @json_serializable_dataclass(frozen=True)
 class FloquetPhasedFSimCalibrationOptions(PhasedFSimCalibrationOptions):
     """Options specific to Floquet PhasedFSimCalibration.
@@ -732,8 +779,36 @@ def _parse_characterized_angles(
     return dict(records)
 
 
+@json_serializable_dataclass(frozen=True)
+class LocalXEBPhasedFSimCalibrationRequest(PhasedFSimCalibrationRequest):
+    """PhasedFSim characterization request for local cross entropy benchmarking (XEB) calibration.
+
+    A "Local" request (corresponding to `LocalXEBPhasedFSimCalibrationOptions`) instructs
+    `cirq_google.run_calibrations` to execute XEB analysis locally (not via the quantum
+    engine). As such, `run_calibrations` can work with any `cirq.Sampler`, not just
+    `QuantumEngineSampler`.
+
+    Attributes:
+        options: local-XEB-specific characterization options.
+    """
+
+    options: LocalXEBPhasedFSimCalibrationOptions
+
+    def parse_result(self, result: CalibrationResult) -> PhasedFSimCalibrationResult:
+        raise NotImplementedError('Not applicable for local calibrations')
+
+    def to_calibration_layer(self) -> CalibrationLayer:
+        raise NotImplementedError('Not applicable for local calibrations')
+
+
 @json_serializable_dataclass(frozen=True)
 class XEBPhasedFSimCalibrationRequest(PhasedFSimCalibrationRequest):
+    """PhasedFSim characterization request for cross entropy benchmarking (XEB) calibration.
+
+    Attributes:
+        options: XEB-specific characterization options.
+    """
+
     options: XEBPhasedFSimCalibrationOptions
 
     def to_calibration_layer(self) -> CalibrationLayer:

cirq-google/cirq_google/calibration/workflow.py

@@ -24,6 +24,7 @@
     Tuple,
     Union,
     cast,
+    TYPE_CHECKING,
 )
 
 from cirq.circuits import Circuit
@@ -55,10 +56,15 @@
     try_convert_sqrt_iswap_to_fsim,
     PhasedFSimCalibrationOptions,
     RequestT,
+    LocalXEBPhasedFSimCalibrationRequest,
 )
+from cirq_google.calibration.xeb_wrapper import run_local_xeb_calibration
 from cirq_google.engine import Engine, QuantumEngineSampler
 from cirq_google.serializable_gate_set import SerializableGateSet
 
+if TYPE_CHECKING:
+    import cirq
+
 _CALIBRATION_IRRELEVANT_GATES = MeasurementGate, SingleQubitGate, WaitGate
 
 
@@ -600,6 +606,53 @@ def _merge_into_calibrations(
     return index
 
 
+def _run_calibrations_via_engine(
+    calibration_requests: Sequence[PhasedFSimCalibrationRequest],
+    engine: Engine,
+    processor_id: str,
+    gate_set: SerializableGateSet,
+    max_layers_per_request: int = 1,
+    progress_func: Optional[Callable[[int, int], None]] = None,
+):
+    """Helper function for run_calibrations.
+
+    This batches and runs calibration requests the normal way: by using engine.run_calibration.
+    This function assumes that all inputs have been validated (by `run_calibrations`).
+    """
+    results = []
+    nested_calibration_layers = [
+        [
+            calibration.to_calibration_layer()
+            for calibration in calibration_requests[offset : offset + max_layers_per_request]
+        ]
+        for offset in range(0, len(calibration_requests), max_layers_per_request)
+    ]
+
+    for cal_layers in nested_calibration_layers:
+        job = engine.run_calibration(cal_layers, processor_id=processor_id, gate_set=gate_set)
+        request_results = job.calibration_results()
+        results += [
+            calibration.parse_result(result)
+            for calibration, result in zip(calibration_requests, request_results)
+        ]
+        if progress_func:
+            progress_func(len(results), len(calibration_requests))
+    return results
+
+
+def _run_local_calibrations_via_sampler(
+    calibration_requests: Sequence[PhasedFSimCalibrationRequest],
+    sampler: 'cirq.Sampler',
+):
+    """Helper function used by `run_calibrations` to run Local calibrations with a Sampler."""
+    return [
+        run_local_xeb_calibration(
+            cast(LocalXEBPhasedFSimCalibrationRequest, calibration_request), sampler
+        )
+        for calibration_request in calibration_requests
+    ]
+
+
 def run_calibrations(
     calibrations: Sequence[PhasedFSimCalibrationRequest],
     sampler: Union[Engine, Sampler],
@@ -638,6 +691,13 @@ def run_calibrations(
     if not calibrations:
         return []
 
+    calibration_request_types = set(type(cr) for cr in calibrations)
+    if len(calibration_request_types) > 1:
+        raise ValueError(
+            f"All calibrations must be of the same type. You gave: {calibration_request_types}"
+        )
+    (calibration_request_type,) = calibration_request_types
+
     if isinstance(sampler, Engine):
         engine: Optional[Engine] = sampler
     elif isinstance(sampler, QuantumEngineSampler):
@@ -649,36 +709,33 @@ def run_calibrations(
 
     if engine is not None:
         if processor_id is None:
-            raise ValueError('processor_id must be provided when running on the engine')
+            raise ValueError('processor_id must be provided.')
         if gate_set is None:
-            raise ValueError('gate_set must be provided when running on the engine')
-
-        results = []
+            raise ValueError('gate_set must be provided.')
+
+        if calibration_request_type == LocalXEBPhasedFSimCalibrationRequest:
+            sampler = engine.sampler(processor_id=processor_id, gate_set=gate_set)
+            return _run_local_calibrations_via_sampler(calibrations, sampler)
+
+        return _run_calibrations_via_engine(
+            calibrations,
+            engine,
+            processor_id,
+            gate_set,
+            max_layers_per_request,
+            progress_func,
+        )
 
-        requests = [
-            [
-                calibration.to_calibration_layer()
-                for calibration in calibrations[offset : offset + max_layers_per_request]
-            ]
-            for offset in range(0, len(calibrations), max_layers_per_request)
-        ]
+    if calibration_request_type == LocalXEBPhasedFSimCalibrationRequest:
+        return _run_local_calibrations_via_sampler(calibrations, sampler=cast(Sampler, sampler))
 
-        for request in requests:
-            job = engine.run_calibration(request, processor_id=processor_id, gate_set=gate_set)
-            request_results = job.calibration_results()
-            results += [
-                calibration.parse_result(result)
-                for calibration, result in zip(calibrations, request_results)
-            ]
-            if progress_func:
-                progress_func(len(results), len(calibrations))
-
-    elif isinstance(sampler, PhasedFSimEngineSimulator):
-        results = sampler.get_calibrations(calibrations)
-    else:
-        raise ValueError(f'Unsupported sampler type {type(sampler)}')
+    if isinstance(sampler, PhasedFSimEngineSimulator):
+        return sampler.get_calibrations(calibrations)
 
-    return results
+    raise ValueError(
+        f'Unsupported sampler/request combination: Sampler {sampler} cannot run '
+        f'calibration request of type {calibration_request_type}'
+    )
 
 
 def make_zeta_chi_gamma_compensation_for_moments(
@@ -1014,10 +1071,10 @@ def run_zeta_chi_gamma_compensation_for_moments(
         circuit, options, gates_translator, merge_subsets=merge_subsets
     )
     characterizations = run_calibrations(
-        requests,
-        sampler,
-        processor_id,
-        gate_set,
+        calibrations=requests,
+        sampler=sampler,
+        processor_id=processor_id,
+        gate_set=gate_set,
         max_layers_per_request=max_layers_per_request,
         progress_func=progress_func,
     )