Parallel XEB: Add option to specify pairs (#6787)

Co-authored-by: Noureldin <noureldinyosri@google.com>

Author: eliottrosenberg

cirq-core/cirq/experiments/two_qubit_xeb.py

@@ -53,6 +53,46 @@ def _manhattan_distance(qubit1: 'cirq.GridQubit', qubit2: 'cirq.GridQubit') -> i
     return abs(qubit1.row - qubit2.row) + abs(qubit1.col - qubit2.col)
 
 
+def qubits_and_pairs(
+    sampler: 'cirq.Sampler',
+    qubits: Optional[Sequence['cirq.GridQubit']] = None,
+    pairs: Optional[Sequence[tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
+) -> Tuple[Sequence['cirq.GridQubit'], Sequence[tuple['cirq.GridQubit', 'cirq.GridQubit']]]:
+    """Extract qubits and pairs from sampler.
+
+
+    If qubits are not provided, then they are extracted from the pairs (if given) or the sampler.
+    If pairs are not provided then all pairs of adjacent qubits are used.
+
+    Args:
+        sampler: The quantum engine or simulator to run the circuits.
+        qubits: Optional list of qubits.
+        pairs: Optional list of pair to use.
+
+    Returns:
+        - Qubits to use.
+        - Pairs of qubits to use.
+
+    Raises:
+        ValueError: If qubits are not specified and can't be deduced from other arguments.
+    """
+    if qubits is None:
+        if pairs is None:
+            qubits = _grid_qubits_for_sampler(sampler)
+            if qubits is None:
+                raise ValueError("Couldn't determine qubits from sampler. Please specify them.")
+        else:
+            qubits_set = set(itertools.chain(*pairs))
+            qubits = list(qubits_set)
+
+    if pairs is None:
+        pairs = [
+            pair for pair in itertools.combinations(qubits, 2) if _manhattan_distance(*pair) == 1
+        ]
+
+    return qubits, pairs
+
+
 @dataclass(frozen=True)
 class TwoQubitXEBResult:
     """Results from an XEB experiment."""
@@ -359,6 +399,7 @@ def parallel_xeb_workflow(
     cycle_depths: Sequence[int] = (5, 25, 50, 100, 200, 300),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     ax: Optional[plt.Axes] = None,
+    pairs: Optional[Sequence[tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
     pool: Optional['multiprocessing.pool.Pool'] = None,
     **plot_kwargs,
 ) -> Tuple[pd.DataFrame, Sequence['cirq.Circuit'], pd.DataFrame]:
@@ -375,6 +416,7 @@ def parallel_xeb_workflow(
         random_state: The random state to use.
         ax: the plt.Axes to plot the device layout on. If not given,
             no plot is created.
+        pairs: Pairs to use. If not specified, use all pairs between adjacent qubits.
         pool: An optional multiprocessing pool.
         **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
 
@@ -391,14 +433,8 @@ def parallel_xeb_workflow(
     """
     rs = value.parse_random_state(random_state)
 
-    if qubits is None:
-        qubits = _grid_qubits_for_sampler(sampler)
-        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
-    )
+    qubits, pairs = qubits_and_pairs(sampler, qubits, pairs)
+    graph = nx.Graph(pairs)
 
     if ax is not None:
         nx.draw_networkx(graph, pos={q: (q.row, q.col) for q in qubits}, ax=ax)
@@ -445,6 +481,7 @@ def parallel_two_qubit_xeb(
     cycle_depths: Sequence[int] = (5, 25, 50, 100, 200, 300),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     ax: Optional[plt.Axes] = None,
+    pairs: Optional[Sequence[tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
     **plot_kwargs,
 ) -> TwoQubitXEBResult:
     """A convenience method that runs the full XEB workflow.
@@ -460,6 +497,7 @@ def parallel_two_qubit_xeb(
         random_state: The random state to use.
         ax: the plt.Axes to plot the device layout on. If not given,
             no plot is created.
+        pairs: Pairs to use. If not specified, use all pairs between adjacent qubits.
         **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
     Returns:
         A TwoQubitXEBResult object representing the results of the experiment.
@@ -469,6 +507,7 @@ def parallel_two_qubit_xeb(
     fids, *_ = parallel_xeb_workflow(
         sampler=sampler,
         qubits=qubits,
+        pairs=pairs,
         entangling_gate=entangling_gate,
         n_repetitions=n_repetitions,
         n_combinations=n_combinations,
@@ -493,6 +532,7 @@ def run_rb_and_xeb(
     depths_xeb: Sequence[int] = (5, 25, 50, 100, 200, 300),
     xeb_combinations: int = 10,
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    pairs: Optional[Sequence[tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
 ) -> InferredXEBResult:
     """A convenience method that runs both RB and XEB workflows.
 
@@ -506,6 +546,7 @@ def run_rb_and_xeb(
         depths_xeb: The cycle depths to use for XEB.
         xeb_combinations: The number of combinations to generate for XEB.
         random_state: The random state to use.
+        pairs: Pairs to use. If not specified, use all pairs between adjacent qubits.
 
     Returns:
         An InferredXEBResult object representing the results of the experiment.
@@ -514,10 +555,7 @@ def run_rb_and_xeb(
         ValueError: If qubits are not specified and the sampler has no device.
     """
 
-    if qubits is None:
-        qubits = _grid_qubits_for_sampler(sampler)
-        if qubits is None:
-            raise ValueError("Couldn't determine qubits from sampler. Please specify them.")
+    qubits, pairs = qubits_and_pairs(sampler, qubits, pairs)
 
     rb = parallel_single_qubit_randomized_benchmarking(
         sampler=sampler,
@@ -530,6 +568,7 @@ def run_rb_and_xeb(
     xeb = parallel_two_qubit_xeb(
         sampler=sampler,
         qubits=qubits,
+        pairs=pairs,
         entangling_gate=entangling_gate,
         n_repetitions=repetitions,
         n_circuits=num_circuits,

cirq-core/cirq/experiments/two_qubit_xeb_test.py

@@ -261,26 +261,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,