1# Copyright 2021 The Cirq Developers
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14import contextlib
15import multiprocessing
16import multiprocessing.pool
17from typing import Optional, Union, Iterator
18
19import cirq
20import cirq.experiments.random_quantum_circuit_generation as rqcg
21import cirq.experiments.xeb_fitting as xebf
22import cirq.experiments.xeb_sampling as xebsamp
23from cirq_google.calibration.phased_fsim import (
24    PhasedFSimCalibrationResult,
25    PhasedFSimCharacterization,
26    LocalXEBPhasedFSimCalibrationRequest,
27    LocalXEBPhasedFSimCalibrationOptions,
28)
29
30
31@contextlib.contextmanager
32def _maybe_multiprocessing_pool(
33    n_processes: Optional[int] = None,
34) -> Iterator[Union['multiprocessing.pool.Pool', None]]:
35    """Yield a multiprocessing.Pool as a context manager, unless n_processes=1; then yield None,
36    which should disable multiprocessing in XEB apis."""
37    if n_processes == 1:
38        yield None
39        return
40
41    with multiprocessing.get_context('spawn').Pool(processes=n_processes) as pool:
42        yield pool
43
44
45def run_local_xeb_calibration(
46    calibration: LocalXEBPhasedFSimCalibrationRequest,
47    sampler: cirq.Sampler,
48) -> PhasedFSimCalibrationResult:
49    """Run a calibration request using `cirq.experiments` XEB utilities and a sampler rather
50    than `Engine.run_calibrations`.
51
52    Args:
53        calibration: A LocalXEBPhasedFSimCalibration request describing the XEB characterization
54            to carry out.
55        sampler: A sampler to execute circuits.
56    """
57    options: LocalXEBPhasedFSimCalibrationOptions = calibration.options
58    circuit = cirq.Circuit([calibration.gate.on(*pair) for pair in calibration.pairs])
59
60    # 2. Set up XEB experiment
61    cycle_depths = options.cycle_depths
62    circuits = rqcg.generate_library_of_2q_circuits(
63        n_library_circuits=options.n_library_circuits,
64        two_qubit_gate=calibration.gate,
65        max_cycle_depth=max(cycle_depths),
66    )
67    combs_by_layer = rqcg.get_random_combinations_for_layer_circuit(
68        n_library_circuits=len(circuits),
69        n_combinations=options.n_combinations,
70        layer_circuit=circuit,
71    )
72
73    # 3. Sample data
74    sampled_df = xebsamp.sample_2q_xeb_circuits(
75        sampler=sampler,
76        circuits=circuits,
77        cycle_depths=cycle_depths,
78        combinations_by_layer=combs_by_layer,
79    )
80
81    # 4. Initial fidelities
82    # initial_fids = xebf.benchmark_2q_xeb_fidelities(
83    #     sampled_df=sampled_df,
84    #     circuits=circuits,
85    #     cycle_depths=cycle_depths,
86    # )
87
88    # 5. Characterize by fitting angles.
89    if options.fsim_options.defaults_set():
90        fsim_options = options.fsim_options
91    else:
92        fsim_options = options.fsim_options.with_defaults_from_gate(calibration.gate)
93
94    pcircuits = [xebf.parameterize_circuit(circuit, fsim_options) for circuit in circuits]
95    fatol = options.fatol if options.fatol is not None else 5e-3
96    xatol = options.xatol if options.xatol is not None else 5e-3
97    with _maybe_multiprocessing_pool(n_processes=options.n_processes) as pool:
98        char_results = xebf.characterize_phased_fsim_parameters_with_xeb_by_pair(
99            sampled_df=sampled_df,
100            parameterized_circuits=pcircuits,
101            cycle_depths=cycle_depths,
102            options=fsim_options,
103            pool=pool,
104            fatol=fatol,
105            xatol=xatol,
106        )
107
108    return PhasedFSimCalibrationResult(
109        parameters={
110            pair: PhasedFSimCharacterization(**param_dict)
111            for pair, param_dict in char_results.final_params.items()
112        },
113        gate=calibration.gate,
114        options=options,
115    )
116