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

cirq/experiments/readout_confusion_matrix.py

@@ -15,11 +15,13 @@
 """Utilities to compute readout confusion matrix and use it for readout error mitigation."""
 
 import time
-from typing import Any, Dict, Union, Sequence, List, Tuple, TYPE_CHECKING, Optional, cast
-import sympy
+from typing import Any, cast, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+
 import numpy as np
 import scipy.optimize
-from cirq import circuits, ops, vis, study
+import sympy
+
+from cirq import circuits, ops, study, vis
 from cirq._compat import proper_repr
 
 if TYPE_CHECKING:
@@ -31,7 +33,7 @@ class TensoredConfusionMatrices:
 
     The confusion matrix (CM) for one qubit is:
 
-        [ Pr(0|0) Pr(1|0) ]
+        [ Pr(0|0) Pr(0|1) ]
         [ Pr(1|0) Pr(1|1) ]
 
     where Pr(i | j) = Probability of observing state "i" given state "j" was prepared.
@@ -71,7 +73,7 @@ class TensoredConfusionMatrices:
             confusion_matrices: Sequence of confusion matrices, computed for qubit patterns present
                                 in `measure_qubits`. A single confusion matrix is also accepted.
             measure_qubits: Sequence of smaller qubit patterns, for which the confusion matrices
-                            were computed. A single qubit pattern is also accepted. Note that the
+                            were computed. A single qubit pattern is also accepted. Note that
                             each qubit pattern is a sequence of qubits used to label the axes of
                             the corresponding confusion matrix.
             repetitions:    The number of repetitions that were used to estimate the confusion
@@ -173,13 +175,12 @@ class TensoredConfusionMatrices:
         return in_vars + out_vars
 
     def _confusion_matrix(self, qubits: Sequence['cirq.Qid']) -> np.ndarray:
-        ein_input = []
+        ein_input: List[np.ndarray | List[int]] = []
         for qs, cm in zip(self.measure_qubits, self.confusion_matrices):
             ein_input.extend([cm.reshape((2, 2) * len(qs)), self._get_vars(qs)])
         ein_out = self._get_vars(qubits)
 
-        # TODO(#5757): remove type ignore when numpy has proper override signature.
-        ret = np.einsum(*ein_input, ein_out).reshape((2 ** len(qubits),) * 2)  # type: ignore
+        ret = np.einsum(*ein_input, ein_out).reshape((2 ** len(qubits),) * 2)
         return ret / ret.sum(axis=1)
 
     def confusion_matrix(self, qubits: Optional[Sequence['cirq.Qid']] = None) -> np.ndarray:
@@ -298,6 +299,70 @@ class TensoredConfusionMatrices:
             )  # pragma: no cover
         return res.x
 
+    def readout_mitigation_pauli_uncorrelated(
+        self, qubits: Sequence['cirq.Qid'], measured_bitstrings: np.ndarray
+    ) -> tuple[float, float]:
+        r"""Uncorrelated readout error mitigation for a multi-qubit Pauli operator.
+
+        This function scalably performs readout error mitigation on an arbitrary-length Pauli
+        operator. It is a reimplementation of https://github.com/eliottrosenberg/correlated_SPAM
+        but specialized to the case in which readout is uncorrelated. We require that the confusion
+        matrix is a tensor product of single-qubit confusion matrices. We then invert the confusion
+        matrix by inverting each of the $C^{(q)}$ Then, in a bit-by-bit fashion, we apply the
+        inverses of the single-site confusion matrices to the bits of the measured bitstring,
+        contract them with the single-site Pauli operator, and take the product over all of the
+        bits. This could be generalized to tensor product spaces that are larger than single qubits,
+        but the essential simplification is that each tensor product space is small, so that none of
+        the response matrices is exponentially large.
+
+        This can result in mitigated Pauli operators that are not in the range [-1, 1], but if
+        the readout error is indeed uncorrelated and well-characterized, then it should converge
+        to being within this range. Results are improved both by a more precise characterization
+        of the response matrices (whose statistical uncertainty is not accounted for in the error
+        propagation here) and by increasing the number of measured bitstrings.
+
+        Args:
+            qubits: The qubits on which the Pauli operator acts.
+            measured_bitstrings: The experimentally measured bitstrings in the eigenbasis of the
+                Pauli operator. measured_bitstrings[i,j] is the ith bitstring, qubit j.
+
+        Returns:
+            The error-mitigated expectation value of the Pauli operator and its statistical
+            uncertainty (not including the uncertainty in the confusion matrices for now).
+
+        Raises:
+            NotImplementedError: If the confusion matrix is not a tensor product of single-qubit
+                                 confusion matrices for all of `qubits`.
+        """
+
+        # in case given as an array of bools, convert to an array of ints:
+        measured_bitstrings = measured_bitstrings.astype(int)
+
+        # get all of the confusion matrices
+        cm_all = []
+        for qubit in qubits:
+            try:
+                idx = self.measure_qubits.index((qubit,))
+            except:  # pragma: no cover
+                raise NotImplementedError(
+                    "The response matrix must be a tensor product of single-qu"
+                    + f"bit response matrices, including that of qubit {qubit}."
+                )
+            cm_all.append(self.confusion_matrices[idx])
+
+        # get the correction matrices, assuming uncorrelated readout:
+        cminv_all = [np.linalg.inv(cm) for cm in cm_all]
+
+        # next, contract them with the single-qubit Pauli operators:
+        z = np.array([1, -1])
+        z_cminv_all = np.array([z @ cminv for cminv in cminv_all])
+
+        # finally, mitigate each bitstring:
+        z_mit_all_shots = np.prod(np.einsum("iji->ij", z_cminv_all[:, measured_bitstrings]), axis=0)
+
+        # return mean and statistical uncertainty:
+        return np.mean(z_mit_all_shots), np.std(z_mit_all_shots) / np.sqrt(len(measured_bitstrings))
+
     def __repr__(self) -> str:
         return (
             f"cirq.TensoredConfusionMatrices("

cirq/experiments/readout_confusion_matrix_test.py

@@ -13,12 +13,40 @@
 # limitations under the License.
 
 import numpy as np
-import cirq
 import pytest
 
+import cirq
+from cirq.experiments.readout_confusion_matrix import TensoredConfusionMatrices
 from cirq.experiments.single_qubit_readout_calibration_test import NoisySingleQubitReadoutSampler
 
 
+def add_readout_error(
+    measurements: np.ndarray,
+    zero_errors: np.ndarray,
+    one_errors: np.ndarray,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Add readout errors to measured (or simulated) bitstrings.
+
+    Args:
+        measurements: The bitstrings to which we will add readout errors. measurements[i,j] is the
+                      ith bitstring, qubit j.
+        zero_errors: zero_errors[i] is the probability of a 0->1 readout error on qubit i.
+        one_errors: one_errors[i] is the probability of a 1->0 readout error on qubit i.
+        rng: The pseudorandom number generator to use.
+
+    Returns:
+        New measurements but with readout errors added.
+    """
+    num_bitstrs, n = measurements.shape
+    assert len(zero_errors) == len(one_errors) == n
+    # compute the probability that each bit is 1 after adding readout errors:
+    p1 = measurements * (1 - one_errors) + (1 - measurements) * zero_errors
+    r = rng.random((num_bitstrs, n))
+    noisy_measurements = r < p1
+    return noisy_measurements.astype(int)
+
+
 def get_expected_cm(num_qubits: int, p0: float, p1: float):
     expected_cm = np.zeros((2**num_qubits,) * 2)
     for i in range(2**num_qubits):
@@ -159,3 +187,78 @@ def test_readout_confusion_matrix_repr_and_equality():
     eq.add_equality_group(a, a)
     eq.add_equality_group(b, b)
     eq.add_equality_group(c, c)
+
+
+def _sample_ghz(n: int, repetitions: int, rng: np.random.Generator) -> np.ndarray:
+    """Sample a GHZ state in the z basis.
+    Args:
+        n: The number of qubits.
+        repetitions: The number of repetitions.
+        rng: The pseudorandom number generator to use.
+    Returns:
+        An array of the measurement outcomes.
+    """
+    return np.tile(rng.integers(0, 2, size=repetitions), (n, 1)).T
+
+
+def _add_noise_and_mitigate_ghz(
+    n: int,
+    repetitions: int,
+    zero_errors: np.ndarray,
+    one_errors: np.ndarray,
+    rng: np.random.Generator | None = None,
+) -> tuple[float, float, float, float]:
+    """Add readout error to GHZ-like bitstrings and measure <ZZZ...> with and
+    without readout error mitigation.
+    Args:
+        n: The number of qubits.
+        repetitions: The number of repetitions.
+        zero_errors: zero_errors[i] is the probability of a 0->1 readout error on qubit i.
+        one_errors: one_errors[i] is the probability of a 1->0 readout error on qubit i.
+        rng: The pseudorandom number generator to use.
+    Returns:
+        A tuple of:
+        - The mitigated expectation value of <ZZZ...>
+        - The statistical uncertainty of the previous output
+        - The unmitigated expectation value of <ZZZ...>
+        - The statstical uncertainty of the previous output
+    """
+    if rng is None:
+        rng = np.random.default_rng(0)
+    confusion_matrices = [
+        np.array([[1 - e0, e1], [e0, 1 - e1]]) for e0, e1 in zip(zero_errors, one_errors)
+    ]
+    qubits = cirq.LineQubit.range(n)
+    tcm = TensoredConfusionMatrices(
+        confusion_matrices, [[q] for q in qubits], repetitions=0, timestamp=0.0
+    )
+
+    measurements = _sample_ghz(n, repetitions, rng)
+    noisy_measurements = add_readout_error(measurements, zero_errors, one_errors, rng)
+    # unmitigated:
+    p1 = np.mean(np.sum(noisy_measurements, axis=1) % 2)
+    z = 1 - 2 * np.mean(p1)
+    dz = 2 * np.sqrt(p1 * (1 - p1) / repetitions)
+    # return mitigated and unmitigated:
+    return (*tcm.readout_mitigation_pauli_uncorrelated(qubits, noisy_measurements), z, dz)
+
+
+def test_uncorrelated_readout_mitigation_pauli():
+    n_all = np.arange(2, 35)
+    z_all_mit = []
+    dz_all_mit = []
+    z_all_raw = []
+    dz_all_raw = []
+    repetitions = 10_000
+    for n in n_all:
+        e0 = np.ones(n) * 0.005
+        e1 = np.ones(n) * 0.03
+        z_mit, dz_mit, z_raw, dz_raw = _add_noise_and_mitigate_ghz(n, repetitions, e0, e1)
+        z_all_mit.append(z_mit)
+        dz_all_mit.append(dz_mit)
+        z_all_raw.append(z_raw)
+        dz_all_raw.append(dz_raw)
+
+    for n, z, dz in zip(n_all, z_all_mit, dz_all_mit):
+        ideal = 1.0 if n % 2 == 0 else 0.0
+        assert np.isclose(z, ideal, atol=4 * dz)

cirq/experiments/single_qubit_readout_calibration.py

@@ -14,15 +14,16 @@
 """Single qubit readout experiments using parallel or isolated statistics."""
 import dataclasses
 import time
-from typing import Any, Dict, Iterable, List, Optional, TYPE_CHECKING
+from typing import Any, cast, Dict, Iterable, List, Optional, TYPE_CHECKING
 
-import sympy
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
+import sympy
+
 import cirq.vis.heatmap as cirq_heatmap
 import cirq.vis.histogram as cirq_histogram
-from cirq.devices import grid_qubit
 from cirq import circuits, ops, study
+from cirq.devices import grid_qubit
 
 if TYPE_CHECKING:
     import cirq
@@ -77,8 +78,9 @@ class SingleQubitReadoutCalibrationResult:
         """
 
         if axs is None:
-            _, axs = plt.subplots(1, 2, dpi=200, facecolor='white', figsize=(12, 4))
-
+            _, axs_v = plt.subplots(1, 2, dpi=200, facecolor='white', figsize=(12, 4))
+            axs_v = cast(np.ndarray, axs_v)
+            axs = cast(tuple[plt.Axes, plt.Axes], (axs_v[0], axs_v[1]))
         else:
             if (
                 not isinstance(axs, (tuple, list, np.ndarray))

cirq/experiments/single_qubit_readout_calibration_test.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from typing import Sequence
-import pytest
 
 import numpy as np
+import pytest
 
 import cirq
 

cirq/experiments/t1_decay_experiment.py

@@ -12,16 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Optional, Sequence, TYPE_CHECKING, cast
-
 import warnings
+from typing import Any, cast, Optional, Sequence, TYPE_CHECKING
+
+import numpy as np
 import pandas as pd
 import sympy
 from matplotlib import pyplot as plt
-import numpy as np
-
 
-from cirq import circuits, ops, study, value, _import
+from cirq import _import, circuits, ops, study, value
 from cirq._compat import proper_repr
 
 if TYPE_CHECKING:

cirq/experiments/t1_decay_experiment_test.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-
 import numpy as np
 import pandas as pd
+import pytest
 import sympy
 
 import cirq

cirq/experiments/t2_decay_experiment.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import enum
-
 from typing import Any, List, Optional, TYPE_CHECKING, Union
 
 import pandas as pd

cirq/experiments/t2_decay_experiment_test.py

@@ -12,9 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-
 import pandas as pd
+import pytest
 import sympy
 
 import cirq

cirq/experiments/two_qubit_xeb.py

@@ -13,31 +13,35 @@
 # limitations under the License.
 
 """Provides functions for running and analyzing two-qubit XEB experiments."""
-from typing import Sequence, TYPE_CHECKING, Optional, Tuple, Dict, cast, Mapping
-
+import functools
+import itertools
 from dataclasses import dataclass
 from types import MappingProxyType
-import itertools
-import functools
+from typing import Any, cast, Dict, Mapping, Optional, Sequence, Tuple, TYPE_CHECKING
 
-from matplotlib import pyplot as plt
 import networkx as nx
 import numpy as np
 import pandas as pd
+from matplotlib import pyplot as plt
 
 from cirq import ops, value, vis
-from cirq.experiments.xeb_sampling import sample_2q_xeb_circuits
-from cirq.experiments.xeb_fitting import benchmark_2q_xeb_fidelities
-from cirq.experiments.xeb_fitting import fit_exponential_decays, exponential_decay
+from cirq._compat import cached_method
 from cirq.experiments import random_quantum_circuit_generation as rqcg
 from cirq.experiments.qubit_characterizations import (
-    ParallelRandomizedBenchmarkingResult,
     parallel_single_qubit_randomized_benchmarking,
+    ParallelRandomizedBenchmarkingResult,
 )
+from cirq.experiments.xeb_fitting import (
+    benchmark_2q_xeb_fidelities,
+    exponential_decay,
+    fit_exponential_decays,
+)
+from cirq.experiments.xeb_sampling import sample_2q_xeb_circuits
 from cirq.qis import noise_utils
-from cirq._compat import cached_method
 
 if TYPE_CHECKING:
+    import multiprocessing
+
     import cirq
 
 
@@ -52,6 +56,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."""
@@ -143,7 +187,9 @@ class TwoQubitXEBResult:
 
     def xeb_fidelity(self, q0: 'cirq.GridQubit', q1: 'cirq.GridQubit') -> float:
         """Return the XEB fidelity of a qubit pair."""
-        return self._record(q0, q1).layer_fid
+        return noise_utils.decay_constant_to_xeb_fidelity(
+            self._record(q0, q1).layer_fid, num_qubits=2
+        )
 
     def xeb_error(self, q0: 'cirq.GridQubit', q1: 'cirq.GridQubit') -> float:
         """Return the XEB error of a qubit pair."""
@@ -177,8 +223,7 @@ class TwoQubitXEBResult:
         """Return the Pauli error of all qubit pairs."""
         return {
             pair: noise_utils.decay_constant_to_pauli_error(
-                noise_utils.xeb_fidelity_to_decay_constant(self.xeb_fidelity(*pair), num_qubits=2),
-                num_qubits=2,
+                self._record(*pair).layer_fid, num_qubits=2
             )
             for pair in self.all_qubit_pairs
         }
@@ -347,19 +392,23 @@ class InferredXEBResult:
         return ax
 
 
-def parallel_two_qubit_xeb(
+def parallel_xeb_workflow(
     sampler: 'cirq.Sampler',
     qubits: Optional[Sequence['cirq.GridQubit']] = None,
     entangling_gate: 'cirq.Gate' = ops.CZ,
     n_repetitions: int = 10**4,
     n_combinations: int = 10,
     n_circuits: int = 20,
-    cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
+    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,
+    batch_size: int = 9,
+    tags: Sequence[Any] = (),
     **plot_kwargs,
-) -> TwoQubitXEBResult:
-    """A convenience method that runs the full XEB workflow.
+) -> Tuple[pd.DataFrame, Sequence['cirq.Circuit'], pd.DataFrame]:
+    """A utility method that runs the full XEB workflow.
 
     Args:
         sampler: The quantum engine or simulator to run the circuits.
@@ -372,24 +421,29 @@ 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.
+        pool: An optional multiprocessing pool.
+        batch_size: We call `run_batch` on the sampler, which can speed up execution in certain
+            environments. The number of (circuit, cycle_depth) tasks to be run in each batch
+            is given by this number.
+        tags: Tags to add to two qubit operations.
         **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
 
     Returns:
-        A TwoQubitXEBResult object representing the results of the experiment.
+        - A DataFrame with columns 'cycle_depth' and 'fidelity'.
+        - The circuits used to perform XEB.
+        - 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.
 
     Raises:
         ValueError: If qubits are not specified and the sampler has no device.
     """
     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)
@@ -397,7 +451,11 @@ def parallel_two_qubit_xeb(
         ax.plot(**plot_kwargs)
 
     circuit_library = rqcg.generate_library_of_2q_circuits(
-        n_library_circuits=n_circuits, two_qubit_gate=entangling_gate, random_state=rs
+        n_library_circuits=n_circuits,
+        two_qubit_gate=entangling_gate,
+        random_state=rs,
+        max_cycle_depth=max(cycle_depths),
+        tags=tags,
     )
 
     combs_by_layer = rqcg.get_random_combinations_for_device(
@@ -414,12 +472,70 @@ def parallel_two_qubit_xeb(
         combinations_by_layer=combs_by_layer,
         shuffle=rs,
         repetitions=n_repetitions,
+        batch_size=batch_size,
     )
 
     fids = benchmark_2q_xeb_fidelities(
-        sampled_df=sampled_df, circuits=circuit_library, cycle_depths=cycle_depths
+        sampled_df=sampled_df, circuits=circuit_library, cycle_depths=cycle_depths, pool=pool
     )
 
+    return fids, circuit_library, sampled_df
+
+
+def parallel_two_qubit_xeb(
+    sampler: 'cirq.Sampler',
+    qubits: Optional[Sequence['cirq.GridQubit']] = None,
+    entangling_gate: 'cirq.Gate' = ops.CZ,
+    n_repetitions: int = 10**4,
+    n_combinations: int = 10,
+    n_circuits: int = 20,
+    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,
+    batch_size: int = 9,
+    tags: Sequence[Any] = (),
+    **plot_kwargs,
+) -> TwoQubitXEBResult:
+    """A convenience method that runs the full XEB workflow.
+
+    Args:
+        sampler: The quantum engine or simulator to run the circuits.
+        qubits: Qubits under test. If none, uses all qubits on the sampler's device.
+        entangling_gate: The entangling gate to use.
+        n_repetitions: The number of repetitions to use.
+        n_combinations: The number of combinations to generate.
+        n_circuits: The number of circuits to generate.
+        cycle_depths: The cycle depths to use.
+        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.
+        batch_size: We call `run_batch` on the sampler, which can speed up execution in certain
+            environments. The number of (circuit, cycle_depth) tasks to be run in each batch
+            is given by this number.
+        tags: Tags to add to two qubit operations.
+        **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
+    Returns:
+        A TwoQubitXEBResult object representing the results of the experiment.
+    Raises:
+        ValueError: If qubits are not specified and the sampler has no device.
+    """
+    fids, *_ = parallel_xeb_workflow(
+        sampler=sampler,
+        qubits=qubits,
+        pairs=pairs,
+        entangling_gate=entangling_gate,
+        n_repetitions=n_repetitions,
+        n_combinations=n_combinations,
+        n_circuits=n_circuits,
+        cycle_depths=cycle_depths,
+        random_state=random_state,
+        ax=ax,
+        batch_size=batch_size,
+        tags=tags,
+        **plot_kwargs,
+    )
     return TwoQubitXEBResult(fit_exponential_decays(fids))
 
 
@@ -432,9 +548,12 @@ def run_rb_and_xeb(
         np.logspace(np.log10(5), np.log10(1000), 5, dtype=int)
     ),
     entangling_gate: 'cirq.Gate' = ops.CZ,
-    depths_xeb: Sequence[int] = tuple(np.arange(3, 100, 20)),
+    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,
+    batch_size: int = 9,
+    tags: Sequence[Any] = (),
 ) -> InferredXEBResult:
     """A convenience method that runs both RB and XEB workflows.
 
@@ -448,6 +567,11 @@ 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.
+        batch_size: We call `run_batch` on the sampler, which can speed up execution in certain
+            environments. The number of (circuit, cycle_depth) tasks to be run in each batch
+            is given by this number.
+        tags: Tags to add to two qubit operations.
 
     Returns:
         An InferredXEBResult object representing the results of the experiment.
@@ -456,10 +580,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,
@@ -472,12 +593,15 @@ 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,
         cycle_depths=depths_xeb,
         n_combinations=xeb_combinations,
         random_state=random_state,
+        batch_size=batch_size,
+        tags=tags,
     )
 
     return InferredXEBResult(rb, xeb)