cirq-core 1.4.0.dev20240529225110__py3-none-any.whl → 1.5.0__py3-none-any.whl

cirq/experiments/two_qubit_xeb_test.py

@@ -12,14 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Wraps Parallel Two Qubit XEB into a few convenience methods."""
-from typing import Optional, Sequence, Dict
-import itertools
 import io
+import itertools
+from typing import Dict, Optional, Sequence
 
 import matplotlib.pyplot as plt
-
-import numpy as np
 import networkx as nx
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -148,10 +147,10 @@ _TEST_RESULT = cirq.experiments.TwoQubitXEBResult(
 @pytest.mark.parametrize(
     'q0,q1,pauli',
     [
-        (cirq.GridQubit(4, 4), cirq.GridQubit(5, 4), 1 / 8),
-        (cirq.GridQubit(5, 3), cirq.GridQubit(6, 3), 1 / 4),
-        (cirq.GridQubit(4, 3), cirq.GridQubit(5, 3), 0.8 + 3 / 40),
-        (cirq.GridQubit(6, 3), cirq.GridQubit(6, 4), 5 / 8),
+        (cirq.GridQubit(4, 4), cirq.GridQubit(5, 4), 0.09374999999999997),
+        (cirq.GridQubit(5, 3), cirq.GridQubit(6, 3), 0.18749999999999994),
+        (cirq.GridQubit(4, 3), cirq.GridQubit(5, 3), 0.65625),
+        (cirq.GridQubit(6, 3), cirq.GridQubit(6, 4), 0.46875),
     ],
 )
 def test_pauli_error(q0: cirq.GridQubit, q1: cirq.GridQubit, pauli: float):
@@ -174,10 +173,10 @@ class MockParallelRandomizedBenchmarkingResult(ParallelRandomizedBenchmarkingRes
 @pytest.mark.parametrize(
     'q0,q1,pauli',
     [
-        (cirq.GridQubit(4, 4), cirq.GridQubit(5, 4), 1 / 8 - 0.03),
-        (cirq.GridQubit(5, 3), cirq.GridQubit(6, 3), 1 / 4 - 0.09),
-        (cirq.GridQubit(4, 3), cirq.GridQubit(5, 3), 0.8 + 3 / 40 - 0.08),
-        (cirq.GridQubit(6, 3), cirq.GridQubit(6, 4), 5 / 8 - 0.13),
+        (cirq.GridQubit(4, 4), cirq.GridQubit(5, 4), 0.09374999999999997 - 0.03),
+        (cirq.GridQubit(5, 3), cirq.GridQubit(6, 3), 0.18749999999999994 - 0.09),
+        (cirq.GridQubit(4, 3), cirq.GridQubit(5, 3), 0.65625 - 0.08),
+        (cirq.GridQubit(6, 3), cirq.GridQubit(6, 4), 0.46875 - 0.13),
     ],
 )
 def test_inferred_pauli_error(q0: cirq.GridQubit, q1: cirq.GridQubit, pauli: float):
@@ -191,10 +190,10 @@ def test_inferred_pauli_error(q0: cirq.GridQubit, q1: cirq.GridQubit, pauli: flo
 @pytest.mark.parametrize(
     'q0,q1,xeb',
     [
-        (cirq.GridQubit(4, 4), cirq.GridQubit(5, 4), 0.076),
-        (cirq.GridQubit(5, 3), cirq.GridQubit(6, 3), 0.128),
-        (cirq.GridQubit(4, 3), cirq.GridQubit(5, 3), 0.636),
-        (cirq.GridQubit(6, 3), cirq.GridQubit(6, 4), 0.396),
+        (cirq.GridQubit(4, 4), cirq.GridQubit(5, 4), 0.050999999999999934),
+        (cirq.GridQubit(5, 3), cirq.GridQubit(6, 3), 0.07799999999999996),
+        (cirq.GridQubit(4, 3), cirq.GridQubit(5, 3), 0.46099999999999997),
+        (cirq.GridQubit(6, 3), cirq.GridQubit(6, 4), 0.2709999999999999),
     ],
 )
 def test_inferred_xeb_error(q0: cirq.GridQubit, q1: cirq.GridQubit, xeb: float):
@@ -224,10 +223,10 @@ def test_inferred_single_qubit_pauli():
 @pytest.mark.parametrize(
     'q0,q1,pauli',
     [
-        (cirq.GridQubit(4, 4), cirq.GridQubit(5, 4), 1 / 8),
-        (cirq.GridQubit(5, 3), cirq.GridQubit(6, 3), 1 / 4),
-        (cirq.GridQubit(4, 3), cirq.GridQubit(5, 3), 0.8 + 3 / 40),
-        (cirq.GridQubit(6, 3), cirq.GridQubit(6, 4), 5 / 8),
+        (cirq.GridQubit(4, 4), cirq.GridQubit(5, 4), 0.09374999999999997),
+        (cirq.GridQubit(5, 3), cirq.GridQubit(6, 3), 0.18749999999999994),
+        (cirq.GridQubit(4, 3), cirq.GridQubit(5, 3), 0.65625),
+        (cirq.GridQubit(6, 3), cirq.GridQubit(6, 4), 0.46875),
     ],
 )
 def test_inferred_two_qubit_pauli(q0: cirq.GridQubit, q1: cirq.GridQubit, pauli: float):
@@ -261,26 +260,43 @@ def test_inferred_plots(ax, target_error, kind):
 
 
 @pytest.mark.parametrize(
-    'sampler,qubits',
+    'sampler,qubits,pairs',
     [
         (
             cirq.DensityMatrixSimulator(
                 seed=0, noise=cirq.ConstantQubitNoiseModel(cirq.amplitude_damp(0.1))
             ),
             cirq.GridQubit.rect(3, 2, 4, 3),
+            None,
+        ),
+        (
+            cirq.DensityMatrixSimulator(
+                seed=0, noise=cirq.ConstantQubitNoiseModel(cirq.amplitude_damp(0.1))
+            ),
+            None,
+            [
+                (cirq.GridQubit(0, 0), cirq.GridQubit(0, 1)),
+                (cirq.GridQubit(0, 0), cirq.GridQubit(1, 0)),
+            ],
         ),
         (
             DensityMatrixSimulatorWithProcessor(
                 seed=0, noise=cirq.ConstantQubitNoiseModel(cirq.amplitude_damp(0.1))
             ),
             None,
+            None,
         ),
     ],
 )
-def test_run_rb_and_xeb(sampler: cirq.Sampler, qubits: Optional[Sequence[cirq.GridQubit]]):
+def test_run_rb_and_xeb(
+    sampler: cirq.Sampler,
+    qubits: Optional[Sequence[cirq.GridQubit]],
+    pairs: Optional[Sequence[tuple[cirq.GridQubit, cirq.GridQubit]]],
+):
     res = cirq.experiments.run_rb_and_xeb(
         sampler=sampler,
         qubits=qubits,
+        pairs=pairs,
         repetitions=100,
         num_clifford_range=tuple(np.arange(3, 10, 1)),
         xeb_combinations=1,

cirq/experiments/xeb_fitting.py

@@ -13,20 +13,23 @@
 # limitations under the License.
 """Estimation of fidelity associated with experimental circuit executions."""
 import dataclasses
-from abc import abstractmethod, ABC
+from abc import ABC, abstractmethod
 from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
 import sympy
-from cirq import circuits, ops, protocols, _import
+
+from cirq import _import, circuits, ops, protocols
 from cirq.experiments.xeb_simulation import simulate_2q_xeb_circuits
 
 if TYPE_CHECKING:
-    import cirq
     import multiprocessing
+
     import scipy.optimize
 
+    import cirq
+
 # We initialize these lazily, otherwise they slow global import speed.
 optimize = _import.LazyLoader("optimize", globals(), "scipy.optimize")
 stats = _import.LazyLoader("stats", globals(), "scipy.stats")
@@ -95,6 +98,11 @@ def benchmark_2q_xeb_fidelities(
     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)
+    # Var[m_u] = Var[sum p(x) * p_sampled(x)]
+    #           = sum p(x)^2 Var[p_sampled(x)]
+    #           = sum p(x)^2 p(x) (1 - p(x))
+    #           = sum p(x)^3 (1 - p(x))
+    df['var_m_u'] = np.sum(pure_probs**3 * (1 - pure_probs), axis=1)
     df['y'] = df['m_u'] - df['u_u']
     df['x'] = df['e_u'] - df['u_u']
     df['numerator'] = df['x'] * df['y']
@@ -103,7 +111,11 @@ def benchmark_2q_xeb_fidelities(
     def per_cycle_depth(df):
         """This function is applied per cycle_depth in the following groupby aggregation."""
         fid_lsq = df['numerator'].sum() / df['denominator'].sum()
-        ret = {'fidelity': fid_lsq}
+        # Note: both df['denominator'] and df['x'] are constants.
+        # Var[f] = Var[df['numerator']] / (sum df['denominator'])^2
+        #           = sum (df['x']^2 * df['var_m_u']) / (sum df['denominator'])^2
+        var_fid = (df['var_m_u'] * df['x'] ** 2).sum() / df['denominator'].sum() ** 2
+        ret = {'fidelity': fid_lsq, 'fidelity_variance': var_fid}
 
         def _try_keep(k):
             """If all the values for a key `k` are the same in this group, we can keep it."""
@@ -126,7 +138,7 @@ def benchmark_2q_xeb_fidelities(
     else:
         groupby_names = ['cycle_depth']
 
-    return df.groupby(groupby_names).apply(per_cycle_depth).reset_index()
+    return df.groupby(groupby_names).apply(per_cycle_depth, include_groups=False).reset_index()
 
 
 class XEBCharacterizationOptions(ABC):
@@ -146,6 +158,69 @@ class XEBCharacterizationOptions(ABC):
         """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 +250,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}.")
 
 
@@ -317,16 +397,21 @@ def SqrtISwapXEBOptions(*args, **kwargs):
 
 
 def parameterize_circuit(
-    circuit: 'cirq.Circuit', options: XEBCharacterizationOptions
+    circuit: 'cirq.Circuit',
+    options: XEBCharacterizationOptions,
+    target_gatefamily: Optional[ops.GateFamily] = None,
 ) -> 'cirq.Circuit':
     """Parameterize PhasedFSim-like gates in a given circuit according to
     `phased_fsim_options`.
     """
+    if isinstance(target_gatefamily, ops.GateFamily):
+        should_parameterize = lambda op: op in target_gatefamily or options.should_parameterize(op)
+    else:
+        should_parameterize = options.should_parameterize
     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
+            gate.on(*op.qubits) if should_parameterize(op) else op for op in moment.operations
         )
         for moment in circuit.moments
     )
@@ -572,6 +657,7 @@ def _fit_exponential_decay(
             fidelities,
             p0=(a_0, layer_fid_0),
             bounds=((0, 0), (1, 1)),
+            maxfev=1000,
         )
     except ValueError:  # pragma: no cover
         return 0, 0, np.inf, np.inf
@@ -580,15 +666,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.
 
@@ -608,13 +685,16 @@ def fit_exponential_decays(fidelities_df: pd.DataFrame) -> pd.DataFrame:
         a, layer_fid, a_std, layer_fid_std = _fit_exponential_decay(
             f1['cycle_depth'], f1['fidelity']
         )
+        fidelity_variance = 0
+        if 'fidelity_variance' in f1:
+            fidelity_variance = f1['fidelity_variance'].values
         record = {
             'a': a,
             'layer_fid': layer_fid,
             'cycle_depths': f1['cycle_depth'].values,
             'fidelities': f1['fidelity'].values,
             'a_std': a_std,
-            'layer_fid_std': layer_fid_std,
+            'layer_fid_std': np.sqrt(layer_fid_std**2 + fidelity_variance),
         }
         return pd.Series(record)
 
@@ -622,7 +702,7 @@ def fit_exponential_decays(fidelities_df: pd.DataFrame) -> pd.DataFrame:
         groupby = ['layer_i', 'pair_i', 'pair']
     else:
         groupby = ['pair']
-    return fidelities_df.groupby(groupby).apply(_per_pair)
+    return fidelities_df.groupby(groupby).apply(_per_pair, include_groups=False)
 
 
 def before_and_after_characterization(

cirq/experiments/xeb_fitting_test.py

@@ -11,9 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from __future__ import annotations
+
 import itertools
 import multiprocessing
-from typing import Iterable
+from typing import Iterable, Iterator
 
 import networkx as nx
 import numpy as np
@@ -23,18 +26,28 @@ import pytest
 import cirq
 import cirq.experiments.random_quantum_circuit_generation as rqcg
 from cirq.experiments.xeb_fitting import (
+    _fit_exponential_decay,
+    before_and_after_characterization,
     benchmark_2q_xeb_fidelities,
-    parameterize_circuit,
-    SqrtISwapXEBOptions,
     characterize_phased_fsim_parameters_with_xeb,
     characterize_phased_fsim_parameters_with_xeb_by_pair,
-    _fit_exponential_decay,
     fit_exponential_decays,
-    before_and_after_characterization,
+    parameterize_circuit,
+    phased_fsim_angles_from_gate,
+    SqrtISwapXEBOptions,
     XEBPhasedFSimCharacterizationOptions,
 )
 from cirq.experiments.xeb_sampling import sample_2q_xeb_circuits
 
+_POOL_NUM_PROCESSES = min(4, multiprocessing.cpu_count())
+
+
+@pytest.fixture
+def pool() -> Iterator[multiprocessing.pool.Pool]:
+    ctx = multiprocessing.get_context()
+    with ctx.Pool(_POOL_NUM_PROCESSES) as pool:
+        yield pool
+
 
 @pytest.fixture(scope='module')
 def circuits_cycle_depths_sampled_df():
@@ -203,7 +216,7 @@ def test_get_initial_simplex():
     assert simplex.shape[1] == len(names)
 
 
-def test_characterize_phased_fsim_parameters_with_xeb():
+def test_characterize_phased_fsim_parameters_with_xeb(pool):
     q0, q1 = cirq.LineQubit.range(2)
     rs = np.random.RandomState(52)
     circuits = [
@@ -228,17 +241,16 @@ def test_characterize_phased_fsim_parameters_with_xeb():
         characterize_phi=False,
     )
     p_circuits = [parameterize_circuit(circuit, options) for circuit in circuits]
-    with multiprocessing.Pool() as pool:
-        result = characterize_phased_fsim_parameters_with_xeb(
-            sampled_df=sampled_df,
-            parameterized_circuits=p_circuits,
-            cycle_depths=cycle_depths,
-            options=options,
-            # speed up with looser tolerances:
-            fatol=1e-2,
-            xatol=1e-2,
-            pool=pool,
-        )
+    result = characterize_phased_fsim_parameters_with_xeb(
+        sampled_df=sampled_df,
+        parameterized_circuits=p_circuits,
+        cycle_depths=cycle_depths,
+        options=options,
+        # speed up with looser tolerances:
+        fatol=1e-2,
+        xatol=1e-2,
+        pool=pool,
+    )
     opt_res = result.optimization_results[(q0, q1)]
     assert np.abs(opt_res.x[0] + np.pi / 4) < 0.1
     assert np.abs(opt_res.fun) < 0.1  # noiseless simulator
@@ -248,7 +260,7 @@ def test_characterize_phased_fsim_parameters_with_xeb():
 
 
 @pytest.mark.parametrize('use_pool', (True, False))
-def test_parallel_full_workflow(use_pool):
+def test_parallel_full_workflow(request, use_pool):
     circuits = rqcg.generate_library_of_2q_circuits(
         n_library_circuits=5,
         two_qubit_gate=cirq.ISWAP**0.5,
@@ -268,10 +280,8 @@ def test_parallel_full_workflow(use_pool):
         combinations_by_layer=combs,
     )
 
-    if use_pool:
-        pool = multiprocessing.Pool()
-    else:
-        pool = None
+    # avoid starting worker pool if it is not needed
+    pool = request.getfixturevalue("pool") if use_pool else None
 
     fids_df_0 = benchmark_2q_xeb_fidelities(
         sampled_df=sampled_df, circuits=circuits, cycle_depths=cycle_depths, pool=pool
@@ -292,8 +302,6 @@ def test_parallel_full_workflow(use_pool):
         xatol=5e-2,
         pool=pool,
     )
-    if pool is not None:
-        pool.terminate()
 
     assert len(result.optimization_results) == graph.number_of_edges()
     for opt_res in result.optimization_results.values():
@@ -354,7 +362,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 +403,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())

cirq/experiments/xeb_sampling.py

@@ -12,30 +12,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Estimation of fidelity associated with experimental circuit executions."""
-import concurrent
 import os
 import time
 import uuid
-from concurrent.futures.thread import ThreadPoolExecutor
 from dataclasses import dataclass
 from typing import (
+    Any,
     Callable,
+    ContextManager,
+    Dict,
     List,
     Optional,
     Sequence,
+    Set,
     Tuple,
     TYPE_CHECKING,
-    Set,
-    ContextManager,
-    Dict,
-    Any,
 )
 
 import numpy as np
 import pandas as pd
 import tqdm
 
-from cirq import ops, devices, value, protocols
+from cirq import devices, ops, protocols, value
 from cirq.circuits import Circuit, Moment
 from cirq.experiments.random_quantum_circuit_generation import CircuitLibraryCombination
 
@@ -276,18 +274,16 @@ def _execute_sample_2q_xeb_tasks_in_batches(
     run_batch = _SampleInBatches(
         sampler=sampler, repetitions=repetitions, combinations_by_layer=combinations_by_layer
     )
-    with ThreadPoolExecutor(max_workers=2) as pool:
-        futures = [pool.submit(run_batch, task_batch) for task_batch in batched_tasks]
 
-        records = []
-        with progress_bar(total=len(batched_tasks) * batch_size) as progress:
-            for future in concurrent.futures.as_completed(futures):
-                new_records = future.result()
-                if dataset_directory is not None:
-                    os.makedirs(f'{dataset_directory}', exist_ok=True)
-                    protocols.to_json(new_records, f'{dataset_directory}/xeb.{uuid.uuid4()}.json')
-                records.extend(new_records)
-                progress.update(batch_size)
+    records = []
+    with progress_bar(total=len(batched_tasks) * batch_size) as progress:
+        for task in batched_tasks:
+            new_records = run_batch(task)
+            if dataset_directory is not None:
+                os.makedirs(f'{dataset_directory}', exist_ok=True)
+                protocols.to_json(new_records, f'{dataset_directory}/xeb.{uuid.uuid4()}.json')
+            records.extend(new_records)
+            progress.update(batch_size)
     return records
 
 

cirq/experiments/xeb_simulation.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Estimation of fidelity associated with experimental circuit executions."""
 from dataclasses import dataclass
-from typing import List, Optional, Sequence, TYPE_CHECKING, Dict, Any
+from typing import Any, Dict, List, Optional, Sequence, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -21,9 +21,10 @@ import pandas as pd
 from cirq import sim, value
 
 if TYPE_CHECKING:
-    import cirq
     import multiprocessing
 
+    import cirq
+
 
 @dataclass(frozen=True)
 class _Simulate2qXEBTask:
@@ -107,7 +108,7 @@ def simulate_2q_xeb_circuits(
         # Need an actual object; not np.random or else multiprocessing will
         # fail to pickle the closure object:
         # https://github.com/quantumlib/Cirq/issues/3717
-        simulator = sim.Simulator(seed=np.random.RandomState())
+        simulator = sim.Simulator(seed=np.random.RandomState(), dtype=np.complex128)
     _simulate_2q_xeb_circuit = _Simulate_2q_XEB_Circuit(simulator=simulator)
 
     tasks = tuple(

cirq/experiments/xeb_simulation_test.py

@@ -11,9 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from __future__ import annotations
+
 import multiprocessing
-from typing import Dict, Any, Optional
-from typing import Sequence
+from typing import Any, Dict, Iterator, Optional, Sequence
 
 import numpy as np
 import pandas as pd
@@ -23,8 +25,17 @@ import cirq
 import cirq.experiments.random_quantum_circuit_generation as rqcg
 from cirq.experiments.xeb_simulation import simulate_2q_xeb_circuits
 
+_POOL_NUM_PROCESSES = min(4, multiprocessing.cpu_count())
+
+
+@pytest.fixture
+def pool() -> Iterator[multiprocessing.pool.Pool]:
+    ctx = multiprocessing.get_context()
+    with ctx.Pool(_POOL_NUM_PROCESSES) as pool:
+        yield pool
 
-def test_simulate_2q_xeb_circuits():
+
+def test_simulate_2q_xeb_circuits(pool):
     q0, q1 = cirq.LineQubit.range(2)
     circuits = [
         rqcg.random_rotations_between_two_qubit_circuit(
@@ -42,8 +53,7 @@ def test_simulate_2q_xeb_circuits():
         assert len(row['pure_probs']) == 4
         assert np.isclose(np.sum(row['pure_probs']), 1)
 
-    with multiprocessing.Pool() as pool:
-        df2 = simulate_2q_xeb_circuits(circuits, cycle_depths, pool=pool)
+    df2 = simulate_2q_xeb_circuits(circuits, cycle_depths, pool=pool)
 
     pd.testing.assert_frame_equal(df, df2)
 
@@ -76,7 +86,7 @@ def _ref_simulate_2q_xeb_circuit(task: Dict[str, Any]):
     tcircuit = circuit[:circuit_depth]
     tcircuit = cirq.resolve_parameters_once(tcircuit, param_resolver=param_resolver)
 
-    pure_sim = cirq.Simulator()
+    pure_sim = cirq.Simulator(dtype=np.complex128)
     psi = pure_sim.simulate(tcircuit)
     psi_vector = psi.final_state_vector
     pure_probs = cirq.state_vector_to_probabilities(psi_vector)
@@ -118,8 +128,8 @@ def _ref_simulate_2q_xeb_circuits(
     return pd.DataFrame(records).set_index(['circuit_i', 'cycle_depth']).sort_index()
 
 
-@pytest.mark.parametrize('multiprocess', (True, False))
-def test_incremental_simulate(multiprocess):
+@pytest.mark.parametrize('use_pool', (True, False))
+def test_incremental_simulate(request, use_pool):
     q0, q1 = cirq.LineQubit.range(2)
     circuits = [
         rqcg.random_rotations_between_two_qubit_circuit(
@@ -129,16 +139,12 @@ def test_incremental_simulate(multiprocess):
     ]
     cycle_depths = np.arange(3, 100, 9, dtype=np.int64)
 
-    if multiprocess:
-        pool = multiprocessing.Pool()
-    else:
-        pool = None
+    # avoid starting worker pool if it is not needed
+    pool = request.getfixturevalue("pool") if use_pool else None
 
     df_ref = _ref_simulate_2q_xeb_circuits(circuits=circuits, cycle_depths=cycle_depths, pool=pool)
 
     df = simulate_2q_xeb_circuits(circuits=circuits, cycle_depths=cycle_depths, pool=pool)
-    if pool is not None:
-        pool.terminate()
     pd.testing.assert_frame_equal(df_ref, df)
 
     # Use below for approximate equality, if e.g. you're using qsim: