[XEB] Support parallel execution

cirq/experiments/__init__.py

@@ -60,6 +60,7 @@
 from cirq.experiments.random_quantum_circuit_generation import (
     GRID_ALIGNED_PATTERN,
     GRID_STAGGERED_PATTERN,
+    HALF_GRID_STAGGERED_PATTERN,
     GridInteractionLayer,
     random_rotations_between_grid_interaction_layers_circuit,
 )

cirq/experiments/fidelity_estimation.py

@@ -25,10 +25,10 @@
     Tuple,
     cast,
     TYPE_CHECKING,
-    Dict,
-    Any,
     Set,
     ContextManager,
+    Dict,
+    Any,
 )
 
 import numpy as np
@@ -37,8 +37,9 @@
 import sympy
 import tqdm
 
-from cirq import ops, protocols, sim
+from cirq import ops, sim, devices, protocols
 from cirq.circuits import Circuit
+from cirq.experiments.random_quantum_circuit_generation import CircuitLibraryCombination
 from cirq.ops import QubitOrder, QubitOrderOrList
 from cirq.sim import final_state_vector
 
@@ -413,16 +414,31 @@ class _Sample2qXEBTask:
     """
 
     cycle_depth: int
-    circuit_i: int
+    layer_i: int
+    combination_i: int
     prepared_circuit: 'cirq.Circuit'
+    combination: List[int]
 
 
 class _SampleInBatches:
-    def __init__(self, sampler: 'cirq.Sampler', repetitions: int):
+    def __init__(
+        self,
+        sampler: 'cirq.Sampler',
+        repetitions: int,
+        combinations_by_layer: List[CircuitLibraryCombination],
+    ):
         """This closure will execute a list of `tasks` with one call to
-        `run_batch` on the provided sampler for a given number of repetitions."""
+        `run_batch` on the provided sampler for a given number of repetitions.
+
+        It also keeps a record of the circuit library combinations in order to
+        back out which qubit pairs correspond to each pair index. We tag
+        our return value with this so it is in the resultant DataFrame, which
+        is very convenient for dealing with the results (but not strictly
+        necessary, as the information could be extracted from (`layer_i`, `pair_i`).
+        """
         self.sampler = sampler
         self.repetitions = repetitions
+        self.combinations_by_layer = combinations_by_layer
 
     def __call__(self, tasks: List[_Sample2qXEBTask]):
         prepared_circuits = [task.prepared_circuit for task in tasks]
@@ -431,16 +447,27 @@ def __call__(self, tasks: List[_Sample2qXEBTask]):
         records = []
         for task, nested_result in zip(tasks, results):
             (result,) = nested_result  # remove nesting due to potential sweeps.
-            sampled_inds = result.data.values[:, 0]
-            sampled_probs = np.bincount(sampled_inds, minlength=2 ** 2) / len(sampled_inds)
-
-            records += [
-                {
-                    'circuit_i': task.circuit_i,
-                    'cycle_depth': task.cycle_depth,
-                    'sampled_probs': sampled_probs,
-                }
-            ]
+            for pair_i, circuit_i in enumerate(task.combination):
+                pair_measurement_key = str(pair_i)
+                q0, q1 = self.combinations_by_layer[task.layer_i].pairs[pair_i]
+                sampled_inds = result.data[pair_measurement_key].values
+                sampled_probs = np.bincount(sampled_inds, minlength=2 ** 2) / len(sampled_inds)
+
+                records += [
+                    {
+                        'circuit_i': circuit_i,
+                        'cycle_depth': task.cycle_depth,
+                        'sampled_probs': sampled_probs,
+                        # Additional metadata to track *how* this circuit
+                        # was zipped and executed.
+                        'layer_i': task.layer_i,
+                        'pair_i': pair_i,
+                        'combination_i': task.combination_i,
+                        'pair_name': f'{q0}-{q1}',
+                        'q0': q0,
+                        'q1': q1,
+                    }
+                ]
         return records
 
 
@@ -456,6 +483,14 @@ def _verify_and_get_two_qubits_from_circuits(circuits: Sequence['cirq.Circuit'])
     return all_qubits_list
 
 
+def _verify_two_line_qubits_from_circuits(circuits: Sequence['cirq.Circuit']):
+    if _verify_and_get_two_qubits_from_circuits(circuits) != devices.LineQubit.range(2):
+        raise ValueError(
+            "`circuits` should be a sequence of circuits each operating "
+            "on LineQubit(0) and LineQubit(1)"
+        )
+
+
 class _NoProgress:
     """Dummy (lack of) tqdm-style progress bar."""
 
@@ -470,17 +505,42 @@ def __enter__(
     def __exit__(self, exc_type, exc_val, exc_tb):
         pass
 
-    def update(self, increment: int):
+    def update(self, n: int = 1):
         pass
 
 
+@dataclass(frozen=True)
+class _ZippedCircuit:
+    """A fully-wide circuit made by zipping together a bunch of two-qubit circuits
+    and its provenance data.
+
+    Args:
+        wide_circuit: The zipped circuit on all pairs
+        pairs: The pairs of qubits operated on in the wide circuit.
+        layer_i: The index of the GridInteractionLayer to which the pairs correspond.
+        combination_i: The row index of the combinations matrix that identifies this
+            particular combination of component narrow circuits.
+        combination: The row of the combinations matrix. Each entry is an index
+            into the (narrow) `circuits` library. Each entry indexes the
+            narrow circuit operating on the corresponding pair in `pairs`.
+    """
+
+    wide_circuit: 'cirq.Circuit'
+    pairs: List[Tuple['cirq.Qid', 'cirq.Qid']]
+    layer_i: int
+    combination_i: int
+    combination: List[int]
+
+
 def sample_2q_xeb_circuits(
     sampler: 'cirq.Sampler',
     circuits: Sequence['cirq.Circuit'],
     cycle_depths: Sequence[int],
+    *,
     repetitions: int = 10_000,
     batch_size: int = 9,
     progress_bar: Optional[Callable[..., ContextManager]] = tqdm.tqdm,
+    combinations_by_layer: Optional[List[CircuitLibraryCombination]] = None,
 ):
     """Sample two-qubit XEB circuits given a sampler.
 
@@ -496,32 +556,97 @@ def sample_2q_xeb_circuits(
             is given by this number.
         progress_bar: A progress context manager following the `tqdm` API or `None` to not report
             progress.
+        combinations_by_layer: Either `None` or the result of
+            `rqcg.get_random_combinations_for_device`. If this is `None`, the circuits specified
+            by `circuits` will be sampled verbatim, resulting in isolated XEB characterization.
+            Otherwise, this contains all the random combinations and metadata required to combine
+            the circuits in `circuits` into wide, parallel-XEB-style circuits for execution.
 
     Returns:
-        A pandas dataframe with index given by ['circuit_i', 'cycle_depth'] and
-        column "sampled_probs".
+        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.
     """
+    # Set up progress reporting
     if progress_bar is None:
         progress_bar = _NoProgress
 
-    q0, q1 = _verify_and_get_two_qubits_from_circuits(circuits)
+    # Shim isolated-XEB as a special case of combination-style parallel XEB.
+    if combinations_by_layer is None:
+        q0, q1 = _verify_and_get_two_qubits_from_circuits(circuits)
+        circuits = [
+            circuit.transform_qubits(
+                lambda q: {q0: devices.LineQubit(0), q1: devices.LineQubit(1)}[q]
+            )
+            for circuit in circuits
+        ]
+        combinations_by_layer = [
+            CircuitLibraryCombination(
+                layer=None,
+                combinations=np.arange(len(circuits))[:, np.newaxis],
+                pairs=[(q0, q1)],
+            )
+        ]
+        one_pair = True
+    else:
+        _verify_two_line_qubits_from_circuits(circuits)
+        one_pair = False
+
+    # Check `combinations_by_layer` is compatible with `circuits`.
+    for layer_combinations in combinations_by_layer:
+        if np.any(layer_combinations.combinations < 0) or np.any(
+            layer_combinations.combinations >= len(circuits)
+        ):
+            raise ValueError("`combinations_by_layer` has invalid indices.")
+
+    # Construct fully-wide "zipped" circuits.
+    zipped_circuits: List[_ZippedCircuit] = []
+    for layer_i, layer_combinations in enumerate(combinations_by_layer):
+        for combination_i, combination in enumerate(layer_combinations.combinations):
+            wide_circuit = Circuit.zip(
+                *(
+                    circuits[i].transform_qubits(lambda q: pair[q.x])
+                    for i, pair in zip(combination, layer_combinations.pairs)
+                )
+            )
+            zipped_circuits.append(
+                _ZippedCircuit(
+                    layer_i=layer_i,
+                    combination_i=combination_i,
+                    wide_circuit=wide_circuit,
+                    pairs=layer_combinations.pairs,
+                    combination=combination.tolist(),
+                )
+            )
+
+    # Construct truncated-with-measurement circuits to run!
     tasks = []
     for cycle_depth in cycle_depths:
-        for circuit_i, circuit in enumerate(circuits):
+        for zipped_circuit in zipped_circuits:
             circuit_depth = cycle_depth * 2 + 1
-            assert circuit_depth <= len(circuit)
-            truncated_circuit = circuit[:circuit_depth]
-            prepared_circuit = truncated_circuit + ops.measure(q0, q1)
+            assert circuit_depth <= len(zipped_circuit.wide_circuit)
+            # Slicing creates a copy, although this isn't documented
+            prepared_circuit = zipped_circuit.wide_circuit[:circuit_depth]
+            for pair_i, pair in enumerate(zipped_circuit.pairs):
+                prepared_circuit += ops.measure(*pair, key=str(pair_i))
             tasks.append(
                 _Sample2qXEBTask(
-                    cycle_depth=cycle_depth, circuit_i=circuit_i, prepared_circuit=prepared_circuit
+                    cycle_depth=cycle_depth,
+                    layer_i=zipped_circuit.layer_i,
+                    combination_i=zipped_circuit.combination_i,
+                    prepared_circuit=prepared_circuit,
+                    combination=zipped_circuit.combination,
                 )
             )
 
+    # Batch and run tasks
     n_tasks = len(tasks)
     batched_tasks = [tasks[i : i + batch_size] for i in range(0, n_tasks, batch_size)]
 
-    run_batch = _SampleInBatches(sampler=sampler, repetitions=repetitions)
+    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]
 
@@ -531,7 +656,11 @@ def sample_2q_xeb_circuits(
                 records += future.result()
                 progress.update(batch_size)
 
-    return pd.DataFrame(records).set_index(['circuit_i', 'cycle_depth'])
+    # Set up the dataframe.
+    df = pd.DataFrame(records).set_index(['circuit_i', 'cycle_depth'])
+    if one_pair:
+        df = df.drop(['layer_i', 'pair_i', 'combination_i'], axis=1)
+    return df
 
 
 def _simulate_2q_xeb_circuit(task: Dict[str, Any]):
@@ -644,10 +773,29 @@ def _summary_stats(row):
     df = df.apply(_summary_stats, axis=1)
 
     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()
-        return pd.Series({'fidelity': fid_lsq})
+        ret = {'fidelity': fid_lsq}
+
+        def _try_keep(k):
+            """If all the values for a key `k` are the same in this group, we can keep it."""
+            if k not in df.columns:
+                return  # coverage: ignore
+            vals = df[k].unique()
+            if len(vals) == 1:
+                ret[k] = vals[0]
+
+        _try_keep('q0')
+        _try_keep('q1')
+        _try_keep('pair_name')
+        return pd.Series(ret)
+
+    if 'pair_i' in df.columns:
+        groupby_names = ['layer_i', 'pair_i', 'cycle_depth']
+    else:
+        groupby_names = ['cycle_depth']
 
-    return df.reset_index().groupby('cycle_depth').apply(per_cycle_depth).reset_index()
+    return df.reset_index().groupby(groupby_names).apply(per_cycle_depth).reset_index()
 
 
 # mypy issue: https://github.com/python/mypy/issues/5374