From 5f884f48cf0c3faa14045c7546e4146c20066f74 Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Tue, 24 Jun 2025 15:48:05 -0700
Subject: [PATCH 1/8] Allow users to run readout benchmarking with sweep

---
 .../cirq/contrib/shuffle_circuits/__init__.py |   1 +
 ...ffle_circuits_with_readout_benchmarking.py | 225 +++++++++++++--
 ...circuits_with_readout_benchmarking_test.py | 269 +++++++++++++-----
 3 files changed, 401 insertions(+), 94 deletions(-)

diff --git a/cirq-core/cirq/contrib/shuffle_circuits/__init__.py b/cirq-core/cirq/contrib/shuffle_circuits/__init__.py
index 7f97f8834a4..ee9b9777d67 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/__init__.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/__init__.py
@@ -15,4 +15,5 @@
 
 from cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking import (
     run_shuffled_with_readout_benchmarking as run_shuffled_with_readout_benchmarking,
+    run_sweep_with_readout_benchmarking as run_sweep_with_readout_benchmarking,
 )
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
index df7fa2bdb92..0ad0027be8b 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
@@ -17,11 +17,12 @@
 from __future__ import annotations
 
 import time
-from typing import TYPE_CHECKING
+from typing import Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
+import sympy
 
-from cirq import circuits, ops, protocols, work
+from cirq import circuits, ops, protocols, study, work
 from cirq.experiments import SingleQubitReadoutCalibrationResult
 
 if TYPE_CHECKING:
@@ -29,9 +30,9 @@
 
 
 def _validate_input(
-    input_circuits: list[circuits.Circuit],
-    circuit_repetitions: int | list[int],
-    rng_or_seed: np.random.Generator | int,
+    input_circuits: Sequence[circuits.Circuit],
+    circuit_repetitions: Union[int, list[int]],
+    rng_or_seed: Union[np.random.Generator, int],
     num_random_bitstrings: int,
     readout_repetitions: int,
 ):
@@ -65,6 +66,22 @@ def _validate_input(
         raise ValueError("Must provide non-zero readout_repetitions for readout calibration.")
 
 
+def _validate_input_with_sweep(
+    input_circuits: Sequence[circuits.Circuit],
+    sweep_params: Sequence[study.Sweepable],
+    circuit_repetitions: Union[int, list[int]],
+    rng_or_seed: Union[np.random.Generator, int],
+    num_random_bitstrings: int,
+    readout_repetitions: int,
+):
+    """Validates the input for the run_sweep_with_readout_benchmarking function."""
+    if not sweep_params:
+        raise ValueError("Sweep parameters must not be empty.")
+    return _validate_input(
+        input_circuits, circuit_repetitions, rng_or_seed, num_random_bitstrings, readout_repetitions
+    )
+
+
 def _generate_readout_calibration_circuits(
     qubits: list[ops.Qid], rng: np.random.Generator, num_random_bitstrings: int
 ) -> tuple[list[circuits.Circuit], np.ndarray]:
@@ -84,6 +101,79 @@ def _generate_readout_calibration_circuits(
     return readout_calibration_circuits, random_bitstrings
 
 
+def _generate_parameterized_readout_calibration_circuit_with_sweep(
+    qubits: list[ops.Qid], rng: np.random.Generator, num_random_bitstrings: int
+) -> tuple[circuits.Circuit, study.Sweepable, np.ndarray]:
+    """Generates a parameterized readout calibration circuit, sweep parameters,
+    and the random bitstrings."""
+    random_bitstrings = rng.integers(0, 2, size=(num_random_bitstrings, len(qubits)))
+
+    exp_symbols = [sympy.Symbol(f'exp_{qubit}') for qubit in qubits]
+    parameterized_readout_calibration_circuit = circuits.Circuit(
+        [ops.X(qubit) ** exp for exp, qubit in zip(exp_symbols, qubits)], ops.M(*qubits, key="m")
+    )
+    sweep_params = []
+    for bitstr in random_bitstrings:
+        sweep_params.append({exp: bit for exp, bit in zip(exp_symbols, bitstr)})
+
+    return parameterized_readout_calibration_circuit, sweep_params, random_bitstrings
+
+
+def _generate_all_readout_calibration_circuits(
+    rng: np.random.Generator,
+    num_random_bitstrings: int,
+    qubits_to_measure: List[List[ops.Qid]],
+    is_sweep: bool,
+) -> Tuple[List[circuits.Circuit], List[np.ndarray], List[study.Sweepable]]:
+    """Generates all readout calibration circuits and random bitstrings."""
+    all_readout_calibration_circuits: list[circuits.Circuit] = []
+    all_random_bitstrings: list[np.ndarray] = []
+    all_readout_sweep_params: list[study.Sweepable] = []
+
+    if num_random_bitstrings <= 0:
+        return all_readout_calibration_circuits, all_random_bitstrings, all_readout_sweep_params
+
+    if not is_sweep:
+        for qubit_group in qubits_to_measure:
+            readout_calibration_circuits, random_bitstrings = (
+                _generate_readout_calibration_circuits(qubit_group, rng, num_random_bitstrings)
+            )
+            all_readout_calibration_circuits.extend(readout_calibration_circuits)
+            all_random_bitstrings.append(random_bitstrings)
+    else:
+        for qubit_group in qubits_to_measure:
+            (parameterized_readout_calibration_circuit, readout_sweep_params, random_bitstrings) = (
+                _generate_parameterized_readout_calibration_circuit_with_sweep(
+                    qubit_group, rng, num_random_bitstrings
+                )
+            )
+            all_readout_calibration_circuits.append(parameterized_readout_calibration_circuit)
+            all_readout_sweep_params.append([readout_sweep_params])
+            all_random_bitstrings.append(random_bitstrings)
+
+    return all_readout_calibration_circuits, all_random_bitstrings, all_readout_sweep_params
+
+
+def _determine_qubits_to_measure(
+    input_circuits: Sequence[circuits.Circuit],
+    qubits: Optional[Union[Sequence[ops.Qid], Sequence[Sequence[ops.Qid]]]],
+) -> List[List[ops.Qid]]:
+    """Determine the qubits to measure based on the input circuits and provided qubits."""
+    # If input qubits is None, extract qubits from input circuits
+    qubits_to_measure: List[List[ops.Qid]] = []
+    if qubits is None:
+        qubits_set: set[ops.Qid] = set()
+        for circuit in input_circuits:
+            qubits_set.update(circuit.all_qubits())
+        qubits_to_measure = [sorted(qubits_set)]
+
+    elif isinstance(qubits[0], ops.Qid):
+        qubits_to_measure = [qubits]  # type: ignore
+    else:
+        qubits_to_measure = qubits  # type: ignore
+    return qubits_to_measure
+
+
 def _shuffle_circuits(
     all_circuits: list[circuits.Circuit], all_repetitions: list[int], rng: np.random.Generator
 ) -> tuple[list[circuits.Circuit], list[int], np.ndarray]:
@@ -97,7 +187,7 @@ def _shuffle_circuits(
 
 
 def _analyze_readout_results(
-    unshuffled_readout_measurements: list[ResultDict],
+    unshuffled_readout_measurements: Union[Sequence[ResultDict], Sequence[study.Result]],
     random_bitstrings: np.ndarray,
     readout_repetitions: int,
     qubits: list[ops.Qid],
@@ -163,8 +253,8 @@ def run_shuffled_with_readout_benchmarking(
     rng_or_seed: np.random.Generator | int,
     num_random_bitstrings: int = 100,
     readout_repetitions: int = 1000,
-    qubits: list[ops.Qid] | list[list[ops.Qid]] | None = None,
-) -> tuple[list[ResultDict], dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]:
+    qubits: Optional[Union[Sequence[ops.Qid], Sequence[Sequence[ops.Qid]]]] = None,
+) -> tuple[Sequence[ResultDict], Dict[Tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]:
     """Run the circuits in a shuffled order with readout error benchmarking.
 
     Args:
@@ -191,35 +281,21 @@ def run_shuffled_with_readout_benchmarking(
         input_circuits, circuit_repetitions, rng_or_seed, num_random_bitstrings, readout_repetitions
     )
 
-    # If input qubits is None, extract qubits from input circuits
-    qubits_to_measure: list[list[ops.Qid]] = []
-    if qubits is None:
-        qubits_set: set[ops.Qid] = set()
-        for circuit in input_circuits:
-            qubits_set.update(circuit.all_qubits())
-        qubits_to_measure = [sorted(qubits_set)]
-    elif isinstance(qubits[0], ops.Qid):
-        qubits_to_measure = [qubits]  # type: ignore
-    else:
-        qubits_to_measure = qubits  # type: ignore
+    qubits_to_measure = _determine_qubits_to_measure(input_circuits, qubits)
 
     # Generate the readout calibration circuits if num_random_bitstrings>0
     # Else all_readout_calibration_circuits and all_random_bitstrings are empty
-    all_readout_calibration_circuits = []
-    all_random_bitstrings = []
-
     rng = (
         rng_or_seed
         if isinstance(rng_or_seed, np.random.Generator)
         else np.random.default_rng(rng_or_seed)
     )
-    if num_random_bitstrings > 0:
-        for qubit_group in qubits_to_measure:
-            readout_calibration_circuits, random_bitstrings = (
-                _generate_readout_calibration_circuits(qubit_group, rng, num_random_bitstrings)
-            )
-            all_readout_calibration_circuits.extend(readout_calibration_circuits)
-            all_random_bitstrings.append(random_bitstrings)
+
+    all_readout_calibration_circuits, all_random_bitstrings, _ = (
+        _generate_all_readout_calibration_circuits(
+            rng, num_random_bitstrings, qubits_to_measure, False
+        )
+    )
 
     # Shuffle the circuits
     if isinstance(circuit_repetitions, int):
@@ -254,3 +330,94 @@ def run_shuffled_with_readout_benchmarking(
         start_idx = end_idx
 
     return unshuffled_input_circuits_measiurements, readout_calibration_results
+
+
+def run_sweep_with_readout_benchmarking(
+    input_circuits: list[circuits.Circuit],
+    sweep_params: Sequence[study.Sweepable],
+    sampler: work.Sampler,
+    circuit_repetitions: Union[int, list[int]],
+    rng_or_seed: Union[np.random.Generator, int],
+    num_random_bitstrings: int = 100,
+    readout_repetitions: int = 1000,
+    qubits: Optional[Union[Sequence[ops.Qid], Sequence[Sequence[ops.Qid]]]] = None,
+) -> tuple[
+    Sequence[Sequence[study.Result]], Dict[Tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]
+]:
+    """Run the sweep circuits with readout error benchmarking (no shuffling).
+    Args:
+        input_circuits: The circuits to run.
+        sweep_params: The sweep parameters for the input circuits.
+        sampler: The sampler to use.
+        circuit_repetitions: The repetitions for `circuits`.
+        rng_or_seed: A random number generator used to generate readout circuits.
+                     Or an integer seed.
+        num_random_bitstrings: The number of random bitstrings for measuring readout.
+            If set to 0, no readout calibration circuits are generated.
+        readout_repetitions: The number of repetitions for each readout bitstring.
+        qubits: The qubits to benchmark readout errors. If None, all qubits in the
+                input_circuits are used. Can be a list of qubits or a list of tuples
+                of qubits.
+    Returns:
+        A tuple containing:
+        - A list of lists of dictionaries with the measurement results.
+        - A dictionary mapping each tuple of qubits to a SingleQubitReadoutCalibrationResult.
+    """
+
+    _validate_input_with_sweep(
+        input_circuits,
+        sweep_params,
+        circuit_repetitions,
+        rng_or_seed,
+        num_random_bitstrings,
+        readout_repetitions,
+    )
+
+    qubits_to_measure = _determine_qubits_to_measure(input_circuits, qubits)
+
+    # Generate the readout calibration circuits (parameterized circuits) and sweep params
+    # if num_random_bitstrings>0
+    # Else all_readout_calibration_circuits and all_random_bitstrings are empty
+    rng = (
+        rng_or_seed
+        if isinstance(rng_or_seed, np.random.Generator)
+        else np.random.default_rng(rng_or_seed)
+    )
+
+    all_readout_calibration_circuits, all_random_bitstrings, all_readout_sweep_params = (
+        _generate_all_readout_calibration_circuits(
+            rng, num_random_bitstrings, qubits_to_measure, True
+        )
+    )
+
+    if isinstance(circuit_repetitions, int):
+        circuit_repetitions = [circuit_repetitions] * len(input_circuits)
+    all_repetitions = circuit_repetitions + [readout_repetitions] * len(
+        all_readout_calibration_circuits
+    )
+
+    # Run the sweep circuits and measure
+    results = sampler.run_batch(
+        input_circuits + all_readout_calibration_circuits,
+        list(sweep_params) + all_readout_sweep_params,
+        repetitions=all_repetitions,
+    )
+
+    timestamp = time.time()
+
+    input_circuits_measurement = results[: len(input_circuits)]
+    readout_measurements = results[len(input_circuits) :]
+
+    # Analyze results
+    readout_calibration_results = {}
+    i = 0
+    for qubit_group, random_bitstrings in zip(qubits_to_measure, all_random_bitstrings):
+        group_measurements = readout_measurements[i]
+        i += 1
+
+        calibration_result = _analyze_readout_results(
+            group_measurements, random_bitstrings, readout_repetitions, qubit_group, timestamp
+        )
+        readout_calibration_results[tuple(qubit_group)] = calibration_result
+
+    return input_circuits_measurement, readout_calibration_results
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
index 46b19794065..95184220547 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
@@ -15,11 +15,17 @@
 from __future__ import annotations
 
 import itertools
+from typing import Sequence
 
 import numpy as np
 import pytest
+import sympy
 
 import cirq
+from cirq.contrib.shuffle_circuits import (
+    run_shuffled_with_readout_benchmarking,
+    run_sweep_with_readout_benchmarking,
+)
 from cirq.experiments import (
     random_quantum_circuit_generation as rqcg,
     SingleQubitReadoutCalibrationResult,
@@ -28,7 +34,7 @@
 from cirq.study import ResultDict
 
 
-def _create_test_circuits(qubits: list[cirq.Qid], n_circuits: int) -> list[cirq.Circuit]:
+def _create_test_circuits(qubits: Sequence[cirq.Qid], n_circuits: int) -> list[cirq.Circuit]:
     """Helper function to generate circuits for testing."""
     if len(qubits) < 2:
         raise ValueError(
@@ -50,32 +56,80 @@ def _create_test_circuits(qubits: list[cirq.Qid], n_circuits: int) -> list[cirq.
     return input_circuits
 
 
-def test_shuffled_circuits_with_readout_benchmarking_errors_no_noise():
-    """Test shuffled circuits with readout benchmarking with no noise from sampler."""
-    qubits = cirq.LineQubit.range(5)
+def _create_test_circuits_with_sweep(
+    qubits: Sequence[cirq.Qid], n_circuits: int
+) -> tuple[list[cirq.Circuit], list[cirq.ParamResolver]]:
+    """Helper function to generate sweep circuits for testing."""
+    if len(qubits) < 2:
+        raise ValueError(
+            "Need at least two qubits to generate two-qubit circuits."
+        )  # pragma: no cover
+    theta_symbol = sympy.Symbol('theta')
+    phi_symbol = sympy.Symbol('phi')
+
+    two_qubit_gates = [cirq.ISWAP, cirq.CNOT]
+
+    input_circuits = []
+    sweep_params: list[cirq.ParamResolver] = []
+    qubit_pairs = list(itertools.combinations(qubits, 2))
+    num_pairs = len(qubit_pairs)
+    for i in range(n_circuits):
+        gate = two_qubit_gates[i % len(two_qubit_gates)]
+        q0, q1 = qubit_pairs[i % num_pairs]
+        circuits = rqcg.generate_library_of_2q_circuits(
+            n_library_circuits=5, two_qubit_gate=gate, q0=q0, q1=q1
+        )
+        for circuit in circuits:
+            circuit += cirq.Circuit(cirq.X(q0) ** theta_symbol, cirq.Y(q1) ** phi_symbol)
+            circuit.append(cirq.measure(*qubits, key="m"))
+            sweep_params.append(cirq.ParamResolver({'theta': 0, 'phi': 1}))
+        input_circuits.extend(circuits)
+
+    return input_circuits, sweep_params
 
-    # Generate random input circuits
-    input_circuits = _create_test_circuits(qubits, 3)
+
+@pytest.mark.parametrize("mode", ["shuffled", "sweep"])
+def test_circuits_with_readout_benchmarking_errors_no_noise(mode: str):
+    """Test shuffled/sweep circuits with readout benchmarking with no noise from sampler."""
+    qubits = cirq.LineQubit.range(5)
 
     sampler = cirq.Simulator()
     circuit_repetitions = 1
     # allow passing a seed
     rng = 123
     readout_repetitions = 1000
+    num_random_bitstrings = 100
+
+    if mode == "shuffled":
+        input_circuits = _create_test_circuits(qubits, 3)
 
-    measurements, readout_calibration_results = (
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        sweep_measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
             input_circuits,
             sampler,
             circuit_repetitions,
             rng,
-            num_random_bitstrings=100,
+            num_random_bitstrings=num_random_bitstrings,
+            readout_repetitions=readout_repetitions,
+        )
+
+        for measurement in sweep_measurements:
+            assert isinstance(measurement, ResultDict)
+    elif mode == "sweep":
+        input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 3)
+
+        measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
+            input_circuits,
+            sweep_params,
+            sampler,
+            circuit_repetitions,
+            rng,
+            num_random_bitstrings=num_random_bitstrings,
             readout_repetitions=readout_repetitions,
         )
-    )
 
-    for measurement in measurements:
-        assert isinstance(measurement, ResultDict)
+        for measurement_group in measurements:  # Treat as a list of lists
+            for single_sweep_measurement in measurement_group:
+                assert isinstance(single_sweep_measurement, ResultDict)
 
     for qlist, readout_calibration_result in readout_calibration_results.items():
         assert isinstance(qlist, tuple)
@@ -88,31 +142,46 @@ def test_shuffled_circuits_with_readout_benchmarking_errors_no_noise():
         assert isinstance(readout_calibration_result.timestamp, float)
 
 
-def test_shuffled_circuits_with_readout_benchmarking_errors_with_noise():
-    """Test shuffled circuits with readout benchmarking with noise from sampler."""
+@pytest.mark.parametrize("mode", ["shuffled", "sweep"])
+def test_circuits_with_readout_benchmarking_errors_with_noise(mode: str):
+    """Test shuffled/sweep circuits with readout benchmarking with noise from sampler."""
     qubits = cirq.LineQubit.range(6)
-
-    # Generate random input circuits
-    input_circuits = _create_test_circuits(qubits, 6)
-
     sampler = NoisySingleQubitReadoutSampler(p0=0.1, p1=0.2, seed=1234)
     circuit_repetitions = 1
     rng = np.random.default_rng()
     readout_repetitions = 1000
+    num_random_bitstrings = 100
+
+    if mode == "shuffled":
+        input_circuits = _create_test_circuits(qubits, 6)
 
-    measurements, readout_calibration_results = (
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
             input_circuits,
             sampler,
             circuit_repetitions,
             rng,
-            num_random_bitstrings=100,
+            num_random_bitstrings=num_random_bitstrings,
+            readout_repetitions=readout_repetitions,
+        )
+
+        for measurement in measurements:
+            assert isinstance(measurement, ResultDict)
+    elif mode == "sweep":
+        input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 6)
+
+        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
+            input_circuits,
+            sweep_params,
+            sampler,
+            circuit_repetitions,
+            rng,
+            num_random_bitstrings=num_random_bitstrings,
             readout_repetitions=readout_repetitions,
         )
-    )
 
-    for measurement in measurements:
-        assert isinstance(measurement, ResultDict)
+        for measurement_group in sweep_measurements:  # Treat as a list of lists
+            for single_sweep_measurement in measurement_group:
+                assert isinstance(single_sweep_measurement, ResultDict)
 
     for qlist, readout_calibration_result in readout_calibration_results.items():
         assert isinstance(qlist, tuple)
@@ -127,33 +196,48 @@ def test_shuffled_circuits_with_readout_benchmarking_errors_with_noise():
         assert isinstance(readout_calibration_result.timestamp, float)
 
 
-def test_shuffled_circuits_with_readout_benchmarking_errors_with_noise_and_input_qubits():
-    """Test shuffled circuits with readout benchmarking with noise from sampler and input qubits."""
+@pytest.mark.parametrize("mode", ["shuffled", "sweep"])
+def test_circuits_with_readout_benchmarking_errors_with_noise_and_input_qubits(mode: str):
+    """Test shuffled/sweep circuits with readout benchmarking with noise from sampler and input qubits."""
     qubits = cirq.LineQubit.range(6)
     readout_qubits = qubits[:4]
 
-    # Generate random input circuits
-    input_circuits = _create_test_circuits(qubits, 6)
-
     sampler = NoisySingleQubitReadoutSampler(p0=0.1, p1=0.3, seed=1234)
     circuit_repetitions = 1
     rng = np.random.default_rng()
     readout_repetitions = 1000
+    num_random_bitstrings = 100
+
+    if mode == "shuffled":
+        input_circuits = _create_test_circuits(qubits, 6)
 
-    measurements, readout_calibration_results = (
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
             input_circuits,
             sampler,
             circuit_repetitions,
             rng,
-            num_random_bitstrings=100,
+            num_random_bitstrings=num_random_bitstrings,
             readout_repetitions=readout_repetitions,
             qubits=readout_qubits,
         )
-    )
+        for measurement in measurements:
+            assert isinstance(measurement, ResultDict)
 
-    for measurement in measurements:
-        assert isinstance(measurement, ResultDict)
+    elif mode == "sweep":
+        input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 6)
+        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
+            input_circuits,
+            sweep_params,
+            sampler,
+            circuit_repetitions,
+            rng,
+            num_random_bitstrings=num_random_bitstrings,
+            readout_repetitions=readout_repetitions,
+            qubits=readout_qubits,
+        )
+        for measurement_group in sweep_measurements:  # Treat as a list of lists
+            for single_sweep_measurement in measurement_group:
+                assert isinstance(single_sweep_measurement, ResultDict)
 
     for qlist, readout_calibration_result in readout_calibration_results.items():
         assert isinstance(qlist, tuple)
@@ -168,24 +252,43 @@ def test_shuffled_circuits_with_readout_benchmarking_errors_with_noise_and_input
         assert isinstance(readout_calibration_result.timestamp, float)
 
 
-def test_shuffled_circuits_with_readout_benchmarking_errors_with_noise_and_lists_input_qubits():
-    """Test shuffled circuits with readout benchmarking with noise from sampler and input qubits."""
+@pytest.mark.parametrize("mode", ["shuffled", "sweep"])
+def test_circuits_with_readout_benchmarking_errors_with_noise_and_lists_input_qubits(mode: str):
+    """Test shuffled/sweep circuits with readout benchmarking with noise from sampler and input qubits."""
     qubits_1 = cirq.LineQubit.range(3)
     qubits_2 = cirq.LineQubit.range(4)
-
     readout_qubits = [qubits_1, qubits_2]
 
-    # Generate random input circuits and append measurements
-    input_circuits = _create_test_circuits(qubits_1, 6) + _create_test_circuits(qubits_2, 4)
-
     sampler = NoisySingleQubitReadoutSampler(p0=0.1, p1=0.3, seed=1234)
     circuit_repetitions = 1
     rng = np.random.default_rng()
     readout_repetitions = 1000
 
-    measurements, readout_calibration_results = (
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+    if mode == "shuffled":
+        input_circuits = _create_test_circuits(qubits_1, 6) + _create_test_circuits(qubits_2, 4)
+
+        measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
+            input_circuits,
+            sampler,
+            circuit_repetitions,
+            rng,
+            num_random_bitstrings=100,
+            readout_repetitions=readout_repetitions,
+            qubits=readout_qubits,
+        )
+
+        for measurement in measurements:
+            assert isinstance(measurement, ResultDict)
+
+    elif mode == "sweep":
+        input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits_1, 6)
+        additional_circuits, additional_sweep_params = _create_test_circuits_with_sweep(qubits_2, 4)
+        input_circuits += additional_circuits
+        sweep_params += additional_sweep_params
+
+        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
             input_circuits,
+            sweep_params,
             sampler,
             circuit_repetitions,
             rng,
@@ -193,10 +296,10 @@ def test_shuffled_circuits_with_readout_benchmarking_errors_with_noise_and_lists
             readout_repetitions=readout_repetitions,
             qubits=readout_qubits,
         )
-    )
 
-    for measurement in measurements:
-        assert isinstance(measurement, ResultDict)
+        for measurement_group in sweep_measurements:  # Treat as a list of lists
+            for single_sweep_measurement in measurement_group:
+                assert isinstance(single_sweep_measurement, ResultDict)
 
     for qlist, readout_calibration_result in readout_calibration_results.items():
         assert isinstance(qlist, tuple)
@@ -211,23 +314,22 @@ def test_shuffled_circuits_with_readout_benchmarking_errors_with_noise_and_lists
         assert isinstance(readout_calibration_result.timestamp, float)
 
 
-def test_can_handle_zero_random_bitstring():
-    """Test shuffled circuits without readout benchmarking."""
+@pytest.mark.parametrize("mode", ["shuffled", "sweep"])
+def test_can_handle_zero_random_bitstring(mode: str):
+    """Test shuffled/sweep circuits without readout benchmarking."""
     qubits_1 = cirq.LineQubit.range(3)
     qubits_2 = cirq.LineQubit.range(4)
-
     readout_qubits = [qubits_1, qubits_2]
 
-    # Generate random input circuits and append measurements
-    input_circuits = _create_test_circuits(qubits_1, 6) + _create_test_circuits(qubits_2, 4)
-
     sampler = NoisySingleQubitReadoutSampler(p0=0.1, p1=0.3, seed=1234)
     circuit_repetitions = 1
     rng = np.random.default_rng()
     readout_repetitions = 1000
 
-    measurements, readout_calibration_results = (
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+    if mode == "shuffled":
+        input_circuits = _create_test_circuits(qubits_1, 6) + _create_test_circuits(qubits_2, 4)
+
+        measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
             input_circuits,
             sampler,
             circuit_repetitions,
@@ -236,10 +338,31 @@ def test_can_handle_zero_random_bitstring():
             readout_repetitions=readout_repetitions,
             qubits=readout_qubits,
         )
-    )
 
-    for measurement in measurements:
-        assert isinstance(measurement, ResultDict)
+        for measurement in measurements:
+            assert isinstance(measurement, ResultDict)
+
+    elif mode == "sweep":
+        input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits_1, 6)
+        additional_circuits, additional_sweep_params = _create_test_circuits_with_sweep(qubits_2, 4)
+        input_circuits += additional_circuits
+        sweep_params += additional_sweep_params
+
+        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
+            input_circuits,
+            sweep_params,
+            sampler,
+            circuit_repetitions,
+            rng,
+            num_random_bitstrings=0,
+            readout_repetitions=readout_repetitions,
+            qubits=readout_qubits,
+        )
+
+        for sweep_measurement in sweep_measurements:
+            for single_sweep_measurement in sweep_measurement:
+                assert isinstance(single_sweep_measurement, ResultDict)
+
     # Check that the readout_calibration_results is empty
     assert len(readout_calibration_results.items()) == 0
 
@@ -247,7 +370,7 @@ def test_can_handle_zero_random_bitstring():
 def test_empty_input_circuits():
     """Test that the input circuits are empty."""
     with pytest.raises(ValueError, match="Input circuits must not be empty."):
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        run_shuffled_with_readout_benchmarking(
             [],
             cirq.ZerosSampler(),
             circuit_repetitions=10,
@@ -261,7 +384,7 @@ def test_non_circuit_input():
     """Test that the input circuits are not of type cirq.Circuit."""
     q = cirq.LineQubit(0)
     with pytest.raises(ValueError, match="Input circuits must be of type cirq.Circuit."):
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        run_shuffled_with_readout_benchmarking(
             [q],
             cirq.ZerosSampler(),
             circuit_repetitions=10,
@@ -276,7 +399,7 @@ def test_no_measurements():
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q))
     with pytest.raises(ValueError, match="Input circuits must have measurements."):
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        run_shuffled_with_readout_benchmarking(
             [circuit],
             cirq.ZerosSampler(),
             circuit_repetitions=10,
@@ -291,7 +414,7 @@ def test_zero_circuit_repetitions():
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
     with pytest.raises(ValueError, match="Must provide non-zero circuit_repetitions."):
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        run_shuffled_with_readout_benchmarking(
             [circuit],
             cirq.ZerosSampler(),
             circuit_repetitions=0,
@@ -309,7 +432,7 @@ def test_mismatch_circuit_repetitions():
         ValueError,
         match="Number of circuit_repetitions must match the number of" + " input circuits.",
     ):
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        run_shuffled_with_readout_benchmarking(
             [circuit],
             cirq.ZerosSampler(),
             circuit_repetitions=[10, 20],
@@ -324,7 +447,7 @@ def test_zero_num_random_bitstrings():
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
     with pytest.raises(ValueError, match="Must provide zero or more num_random_bitstrings."):
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        run_shuffled_with_readout_benchmarking(
             [circuit],
             cirq.ZerosSampler(),
             circuit_repetitions=10,
@@ -341,7 +464,7 @@ def test_zero_readout_repetitions():
     with pytest.raises(
         ValueError, match="Must provide non-zero readout_repetitions for readout" + " calibration."
     ):
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        run_shuffled_with_readout_benchmarking(
             [circuit],
             cirq.ZerosSampler(),
             circuit_repetitions=10,
@@ -356,7 +479,7 @@ def test_rng_type_mismatch():
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
     with pytest.raises(ValueError, match="Must provide a numpy random generator or a seed"):
-        cirq.contrib.shuffle_circuits.run_shuffled_with_readout_benchmarking(
+        run_shuffled_with_readout_benchmarking(
             [circuit],
             cirq.ZerosSampler(),
             circuit_repetitions=10,
@@ -364,3 +487,19 @@ def test_rng_type_mismatch():
             num_random_bitstrings=5,
             readout_repetitions=100,
         )
+
+
+def test_empty_sweep_params():
+    """Test that the sweep params are empty."""
+    q = cirq.LineQubit(5)
+    circuit = cirq.Circuit(cirq.H(q))
+    with pytest.raises(ValueError, match="Sweep parameters must not be empty."):
+        run_sweep_with_readout_benchmarking(
+            [circuit],
+            [],
+            cirq.ZerosSampler(),
+            circuit_repetitions=10,
+            rng_or_seed=np.random.default_rng(456),
+            num_random_bitstrings=5,
+            readout_repetitions=100,
+        )

From 684ffa8be3490b47978913760abdf9e035591c1b Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Tue, 1 Jul 2025 23:09:54 -0700
Subject: [PATCH 2/8] Make changes based on Nour's suggestions

---
 .../cirq/contrib/shuffle_circuits/__init__.py |   1 +
 ...ffle_circuits_with_readout_benchmarking.py | 232 ++++++++++++-----
 ...circuits_with_readout_benchmarking_test.py | 240 ++++++++----------
 3 files changed, 280 insertions(+), 193 deletions(-)

diff --git a/cirq-core/cirq/contrib/shuffle_circuits/__init__.py b/cirq-core/cirq/contrib/shuffle_circuits/__init__.py
index ee9b9777d67..a2b90c4f96b 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/__init__.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/__init__.py
@@ -15,5 +15,6 @@
 
 from cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking import (
     run_shuffled_with_readout_benchmarking as run_shuffled_with_readout_benchmarking,
+    run_shuffled_circuits_with_readout_benchmarking as run_shuffled_circuits_with_readout_benchmarking,
     run_sweep_with_readout_benchmarking as run_sweep_with_readout_benchmarking,
 )
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
index 0ad0027be8b..0056dc39082 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
@@ -19,22 +19,64 @@
 import time
 from typing import Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
+import attrs
 import numpy as np
 import sympy
 
 from cirq import circuits, ops, protocols, study, work
 from cirq.experiments import SingleQubitReadoutCalibrationResult
+from cirq._compat import deprecated
 
 if TYPE_CHECKING:
     from cirq.study import ResultDict
 
 
-def _validate_input(
+@attrs.frozen
+class ReadoutBenchmarkingParams:
+    """Parameters for configuring readout benchmarking.
+
+    Attributes:
+        circuit_repetitions: The repetitions for `circuits`.
+        num_random_bitstrings: The number of random bitstrings for measuring readout.
+            If set to 0, no readout calibration circuits are generated.
+        readout_repetitions: The number of repetitions for each readout bitstring.
+        qubits: The qubits to benchmark readout errors. If None, all qubits in the
+                input_circuits are used. Can be a list of qubits or a list of tuples
+                of qubits.
+    """
+
+    circuit_repetitions: Union[int, list[int]]
+    num_random_bitstrings: int = 100
+    readout_repetitions: int = 1000
+    qubits: Optional[Union[Sequence[ops.Qid], Sequence[Sequence[ops.Qid]]]] = None
+
+    def __attrs_post_init__(self):
+        _validate_benchmarking_setup(
+            self.circuit_repetitions, self.num_random_bitstrings, self.readout_repetitions
+        )
+
+
+def _validate_benchmarking_setup(
+    circuit_repetitions: Union[int, list[int]], num_random_bitstrings: int, readout_repetitions: int
+):
+    # Check circuit_repetitions
+    if isinstance(circuit_repetitions, int):
+        if circuit_repetitions <= 0:
+            raise ValueError("Must provide non-zero circuit_repetitions.")
+
+    # Check num_random_bitstrings is bigger than or equal to 0
+    if num_random_bitstrings < 0:
+        raise ValueError("Must provide zero or more num_random_bitstrings.")
+
+    # Check readout_repetitions is bigger than 0
+    if readout_repetitions <= 0:
+        raise ValueError("Must provide non-zero readout_repetitions for readout calibration.")
+
+
+def _validate_experiment_input(
     input_circuits: Sequence[circuits.Circuit],
     circuit_repetitions: Union[int, list[int]],
     rng_or_seed: Union[np.random.Generator, int],
-    num_random_bitstrings: int,
-    readout_repetitions: int,
 ):
     if not input_circuits:
         raise ValueError("Input circuits must not be empty.")
@@ -46,10 +88,6 @@ def _validate_input(
         if not any(protocols.is_measurement(circuit) for op in circuit.all_operations()):
             raise ValueError("Input circuits must have measurements.")
 
-    # Check circuit_repetitions
-    if isinstance(circuit_repetitions, int):
-        if circuit_repetitions <= 0:
-            raise ValueError("Must provide non-zero circuit_repetitions.")
     if isinstance(circuit_repetitions, list) and len(circuit_repetitions) != len(input_circuits):
         raise ValueError("Number of circuit_repetitions must match the number of input circuits.")
 
@@ -57,33 +95,21 @@ def _validate_input(
     if not isinstance(rng_or_seed, np.random.Generator) and not isinstance(rng_or_seed, int):
         raise ValueError("Must provide a numpy random generator or a seed")
 
-    # Check num_random_bitstrings is bigger than or equal to 0
-    if num_random_bitstrings < 0:
-        raise ValueError("Must provide zero or more num_random_bitstrings.")
 
-    # Check readout_repetitions is bigger than 0
-    if readout_repetitions <= 0:
-        raise ValueError("Must provide non-zero readout_repetitions for readout calibration.")
-
-
-def _validate_input_with_sweep(
+def _validate_experiment_input_with_sweep(
     input_circuits: Sequence[circuits.Circuit],
     sweep_params: Sequence[study.Sweepable],
     circuit_repetitions: Union[int, list[int]],
     rng_or_seed: Union[np.random.Generator, int],
-    num_random_bitstrings: int,
-    readout_repetitions: int,
 ):
     """Validates the input for the run_sweep_with_readout_benchmarking function."""
     if not sweep_params:
         raise ValueError("Sweep parameters must not be empty.")
-    return _validate_input(
-        input_circuits, circuit_repetitions, rng_or_seed, num_random_bitstrings, readout_repetitions
-    )
+    return _validate_experiment_input(input_circuits, circuit_repetitions, rng_or_seed)
 
 
 def _generate_readout_calibration_circuits(
-    qubits: list[ops.Qid], rng: np.random.Generator, num_random_bitstrings: int
+    qubits: list[ops.Qid], num_random_bitstrings: int, rng: np.random.Generator
 ) -> tuple[list[circuits.Circuit], np.ndarray]:
     """Generates the readout calibration circuits with random bitstrings."""
     bit_to_gate = (ops.I, ops.X)
@@ -102,10 +128,27 @@ def _generate_readout_calibration_circuits(
 
 
 def _generate_parameterized_readout_calibration_circuit_with_sweep(
-    qubits: list[ops.Qid], rng: np.random.Generator, num_random_bitstrings: int
+    qubits: list[ops.Qid], num_random_bitstrings: int, rng: np.random.Generator
 ) -> tuple[circuits.Circuit, study.Sweepable, np.ndarray]:
     """Generates a parameterized readout calibration circuit, sweep parameters,
-    and the random bitstrings."""
+    and the random bitstrings.
+
+    The function generates a single cirq.Circuit with parameterized X gates.
+    The function also generates a set of random bitstrings and creates a list
+    of sweep parameters to map the parameters in the circuit to the values in
+    each bitstring, allowing efficient calibration of readout errors of input qubits.
+
+    Args:
+        qubits: The list of qubits to include in the calibration circuit.
+        num_random_bitstrings: The number of random bitstrings to generate for calibration.
+        rng: A numpy random number generator used to generate the random bitstrings.
+
+    Returns:
+        A tuple containing:
+            - The parameterized readout calibration circuit (cirq.Circuit).
+            - A list of parameter sweeps (one for each random bitstring).
+            - The numpy array of generated random bitstrings.
+    """
     random_bitstrings = rng.integers(0, 2, size=(num_random_bitstrings, len(qubits)))
 
     exp_symbols = [sympy.Symbol(f'exp_{qubit}') for qubit in qubits]
@@ -120,10 +163,10 @@ def _generate_parameterized_readout_calibration_circuit_with_sweep(
 
 
 def _generate_all_readout_calibration_circuits(
-    rng: np.random.Generator,
     num_random_bitstrings: int,
     qubits_to_measure: List[List[ops.Qid]],
     is_sweep: bool,
+    rng: np.random.Generator,
 ) -> Tuple[List[circuits.Circuit], List[np.ndarray], List[study.Sweepable]]:
     """Generates all readout calibration circuits and random bitstrings."""
     all_readout_calibration_circuits: list[circuits.Circuit] = []
@@ -136,7 +179,7 @@ def _generate_all_readout_calibration_circuits(
     if not is_sweep:
         for qubit_group in qubits_to_measure:
             readout_calibration_circuits, random_bitstrings = (
-                _generate_readout_calibration_circuits(qubit_group, rng, num_random_bitstrings)
+                _generate_readout_calibration_circuits(qubit_group, num_random_bitstrings, rng)
             )
             all_readout_calibration_circuits.extend(readout_calibration_circuits)
             all_random_bitstrings.append(random_bitstrings)
@@ -144,7 +187,7 @@ def _generate_all_readout_calibration_circuits(
         for qubit_group in qubits_to_measure:
             (parameterized_readout_calibration_circuit, readout_sweep_params, random_bitstrings) = (
                 _generate_parameterized_readout_calibration_circuit_with_sweep(
-                    qubit_group, rng, num_random_bitstrings
+                    qubit_group, num_random_bitstrings, rng
                 )
             )
             all_readout_calibration_circuits.append(parameterized_readout_calibration_circuit)
@@ -162,10 +205,9 @@ def _determine_qubits_to_measure(
     # If input qubits is None, extract qubits from input circuits
     qubits_to_measure: List[List[ops.Qid]] = []
     if qubits is None:
-        qubits_set: set[ops.Qid] = set()
-        for circuit in input_circuits:
-            qubits_set.update(circuit.all_qubits())
-        qubits_to_measure = [sorted(qubits_set)]
+        qubits_to_measure = [
+            sorted(set(q for circuit in input_circuits for q in circuit.all_qubits()))
+        ]
 
     elif isinstance(qubits[0], ops.Qid):
         qubits_to_measure = [qubits]  # type: ignore
@@ -246,6 +288,7 @@ def _analyze_readout_results(
     )
 
 
+@deprecated(deadline="v2.0", fix="Use run_shuffled_circuits_with_readout_benchmarking() instead.")
 def run_shuffled_with_readout_benchmarking(
     input_circuits: list[circuits.Circuit],
     sampler: work.Sampler,
@@ -277,9 +320,8 @@ def run_shuffled_with_readout_benchmarking(
 
     """
 
-    _validate_input(
-        input_circuits, circuit_repetitions, rng_or_seed, num_random_bitstrings, readout_repetitions
-    )
+    _validate_benchmarking_setup(circuit_repetitions, num_random_bitstrings, readout_repetitions)
+    _validate_experiment_input(input_circuits, circuit_repetitions, rng_or_seed)
 
     qubits_to_measure = _determine_qubits_to_measure(input_circuits, qubits)
 
@@ -293,7 +335,7 @@ def run_shuffled_with_readout_benchmarking(
 
     all_readout_calibration_circuits, all_random_bitstrings, _ = (
         _generate_all_readout_calibration_circuits(
-            rng, num_random_bitstrings, qubits_to_measure, False
+            num_random_bitstrings, qubits_to_measure, False, rng
         )
     )
 
@@ -332,48 +374,115 @@ def run_shuffled_with_readout_benchmarking(
     return unshuffled_input_circuits_measiurements, readout_calibration_results
 
 
+def run_shuffled_circuits_with_readout_benchmarking(
+    sampler: work.Sampler,
+    input_circuits: list[circuits.Circuit],
+    parameters: ReadoutBenchmarkingParams,
+    rng_or_seed: np.random.Generator | int,
+) -> tuple[Sequence[ResultDict], Dict[Tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]:
+    """Run the circuits in a shuffled order with readout error benchmarking.
+
+    Args:
+        sampler: The sampler to use.
+        input_circuits: The circuits to run.
+        parameters: The readout benchmarking parameters.
+        rng_or_seed: A random number generator used to generate readout circuits.
+                     Or an integer seed.
+
+    Returns:
+        A tuple containing:
+        - A list of dictionaries with the unshuffled measurement results.
+        - A dictionary mapping each tuple of qubits to a SingleQubitReadoutCalibrationResult.
+
+    """
+
+    _validate_experiment_input(input_circuits, parameters.circuit_repetitions, rng_or_seed)
+
+    qubits_to_measure = _determine_qubits_to_measure(input_circuits, parameters.qubits)
+
+    # Generate the readout calibration circuits if num_random_bitstrings>0
+    # Else all_readout_calibration_circuits and all_random_bitstrings are empty
+    rng = (
+        rng_or_seed
+        if isinstance(rng_or_seed, np.random.Generator)
+        else np.random.default_rng(rng_or_seed)
+    )
+
+    all_readout_calibration_circuits, all_random_bitstrings, _ = (
+        _generate_all_readout_calibration_circuits(
+            parameters.num_random_bitstrings, qubits_to_measure, False, rng
+        )
+    )
+
+    # Shuffle the circuits
+    circuit_repetitions = parameters.circuit_repetitions
+    if isinstance(circuit_repetitions, int):
+        circuit_repetitions = [circuit_repetitions] * len(input_circuits)
+    all_repetitions = circuit_repetitions + [parameters.readout_repetitions] * len(
+        all_readout_calibration_circuits
+    )
+
+    shuffled_circuits, all_repetitions, unshuf_order = _shuffle_circuits(
+        input_circuits + all_readout_calibration_circuits, all_repetitions, rng
+    )
+
+    # Run the shuffled circuits and measure
+    results = sampler.run_batch(shuffled_circuits, repetitions=all_repetitions)
+    timestamp = time.time()
+    shuffled_measurements = [res[0] for res in results]
+    unshuffled_measurements = [shuffled_measurements[i] for i in unshuf_order]
+
+    unshuffled_input_circuits_measiurements = unshuffled_measurements[: len(input_circuits)]
+    unshuffled_readout_measurements = unshuffled_measurements[len(input_circuits) :]
+
+    # Analyze results
+    readout_calibration_results = {}
+    start_idx = 0
+    for qubit_group, random_bitstrings in zip(qubits_to_measure, all_random_bitstrings):
+        end_idx = start_idx + len(random_bitstrings)
+        group_measurements = unshuffled_readout_measurements[start_idx:end_idx]
+        calibration_result = _analyze_readout_results(
+            group_measurements,
+            random_bitstrings,
+            parameters.readout_repetitions,
+            qubit_group,
+            timestamp,
+        )
+        readout_calibration_results[tuple(qubit_group)] = calibration_result
+        start_idx = end_idx
+
+    return unshuffled_input_circuits_measiurements, readout_calibration_results
+
+
 def run_sweep_with_readout_benchmarking(
+    sampler: work.Sampler,
     input_circuits: list[circuits.Circuit],
     sweep_params: Sequence[study.Sweepable],
-    sampler: work.Sampler,
-    circuit_repetitions: Union[int, list[int]],
+    parameters: ReadoutBenchmarkingParams,
     rng_or_seed: Union[np.random.Generator, int],
-    num_random_bitstrings: int = 100,
-    readout_repetitions: int = 1000,
-    qubits: Optional[Union[Sequence[ops.Qid], Sequence[Sequence[ops.Qid]]]] = None,
 ) -> tuple[
     Sequence[Sequence[study.Result]], Dict[Tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]
 ]:
     """Run the sweep circuits with readout error benchmarking (no shuffling).
     Args:
+        sampler: The sampler to use.
         input_circuits: The circuits to run.
         sweep_params: The sweep parameters for the input circuits.
-        sampler: The sampler to use.
-        circuit_repetitions: The repetitions for `circuits`.
+        parameters: The readout benchmarking parameters.
         rng_or_seed: A random number generator used to generate readout circuits.
                      Or an integer seed.
-        num_random_bitstrings: The number of random bitstrings for measuring readout.
-            If set to 0, no readout calibration circuits are generated.
-        readout_repetitions: The number of repetitions for each readout bitstring.
-        qubits: The qubits to benchmark readout errors. If None, all qubits in the
-                input_circuits are used. Can be a list of qubits or a list of tuples
-                of qubits.
+
     Returns:
         A tuple containing:
         - A list of lists of dictionaries with the measurement results.
         - A dictionary mapping each tuple of qubits to a SingleQubitReadoutCalibrationResult.
     """
 
-    _validate_input_with_sweep(
-        input_circuits,
-        sweep_params,
-        circuit_repetitions,
-        rng_or_seed,
-        num_random_bitstrings,
-        readout_repetitions,
+    _validate_experiment_input_with_sweep(
+        input_circuits, sweep_params, parameters.circuit_repetitions, rng_or_seed
     )
 
-    qubits_to_measure = _determine_qubits_to_measure(input_circuits, qubits)
+    qubits_to_measure = _determine_qubits_to_measure(input_circuits, parameters.qubits)
 
     # Generate the readout calibration circuits (parameterized circuits) and sweep params
     # if num_random_bitstrings>0
@@ -386,13 +495,14 @@ def run_sweep_with_readout_benchmarking(
 
     all_readout_calibration_circuits, all_random_bitstrings, all_readout_sweep_params = (
         _generate_all_readout_calibration_circuits(
-            rng, num_random_bitstrings, qubits_to_measure, True
+            parameters.num_random_bitstrings, qubits_to_measure, True, rng
         )
     )
 
+    circuit_repetitions = parameters.circuit_repetitions
     if isinstance(circuit_repetitions, int):
         circuit_repetitions = [circuit_repetitions] * len(input_circuits)
-    all_repetitions = circuit_repetitions + [readout_repetitions] * len(
+    all_repetitions = circuit_repetitions + [parameters.readout_repetitions] * len(
         all_readout_calibration_circuits
     )
 
@@ -416,7 +526,11 @@ def run_sweep_with_readout_benchmarking(
         i += 1
 
         calibration_result = _analyze_readout_results(
-            group_measurements, random_bitstrings, readout_repetitions, qubit_group, timestamp
+            group_measurements,
+            random_bitstrings,
+            parameters.readout_repetitions,
+            qubit_group,
+            timestamp,
         )
         readout_calibration_results[tuple(qubit_group)] = calibration_result
 
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
index 95184220547..958aa9062b4 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
@@ -17,14 +17,19 @@
 import itertools
 from typing import Sequence
 
+import cirq.contrib.shuffle_circuits
 import numpy as np
 import pytest
 import sympy
 
 import cirq
 from cirq.contrib.shuffle_circuits import (
-    run_shuffled_with_readout_benchmarking,
+    run_shuffled_circuits_with_readout_benchmarking,
     run_sweep_with_readout_benchmarking,
+    run_shuffled_with_readout_benchmarking,
+)
+from cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking import (
+    ReadoutBenchmarkingParams,
 )
 from cirq.experiments import (
     random_quantum_circuit_generation as rqcg,
@@ -95,39 +100,36 @@ def test_circuits_with_readout_benchmarking_errors_no_noise(mode: str):
 
     sampler = cirq.Simulator()
     circuit_repetitions = 1
+    num_random_bitstrings = 100
+    readout_repetitions = 1000
+
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=circuit_repetitions,
+        num_random_bitstrings=num_random_bitstrings,
+        readout_repetitions=readout_repetitions,
+        qubits=qubits,
+    )
     # allow passing a seed
     rng = 123
-    readout_repetitions = 1000
-    num_random_bitstrings = 100
 
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits, 3)
 
-        sweep_measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
-            input_circuits,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=num_random_bitstrings,
-            readout_repetitions=readout_repetitions,
+        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
+            sampler, input_circuits, readout_benchmarking_params, rng
         )
 
-        for measurement in sweep_measurements:
+        for measurement in measurements:
             assert isinstance(measurement, ResultDict)
+
     elif mode == "sweep":
         input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 3)
 
-        measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            input_circuits,
-            sweep_params,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=num_random_bitstrings,
-            readout_repetitions=readout_repetitions,
+        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
+            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
         )
 
-        for measurement_group in measurements:  # Treat as a list of lists
+        for measurement_group in sweep_measurements:
             for single_sweep_measurement in measurement_group:
                 assert isinstance(single_sweep_measurement, ResultDict)
 
@@ -152,16 +154,18 @@ def test_circuits_with_readout_benchmarking_errors_with_noise(mode: str):
     readout_repetitions = 1000
     num_random_bitstrings = 100
 
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=circuit_repetitions,
+        num_random_bitstrings=num_random_bitstrings,
+        readout_repetitions=readout_repetitions,
+        qubits=qubits,
+    )
+
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits, 6)
 
-        measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
-            input_circuits,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=num_random_bitstrings,
-            readout_repetitions=readout_repetitions,
+        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
+            sampler, input_circuits, readout_benchmarking_params, rng
         )
 
         for measurement in measurements:
@@ -170,13 +174,7 @@ def test_circuits_with_readout_benchmarking_errors_with_noise(mode: str):
         input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 6)
 
         sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            input_circuits,
-            sweep_params,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=num_random_bitstrings,
-            readout_repetitions=readout_repetitions,
+            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
         )
 
         for measurement_group in sweep_measurements:  # Treat as a list of lists
@@ -208,17 +206,18 @@ def test_circuits_with_readout_benchmarking_errors_with_noise_and_input_qubits(m
     readout_repetitions = 1000
     num_random_bitstrings = 100
 
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=circuit_repetitions,
+        num_random_bitstrings=num_random_bitstrings,
+        readout_repetitions=readout_repetitions,
+        qubits=readout_qubits,
+    )
+
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits, 6)
 
-        measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
-            input_circuits,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=num_random_bitstrings,
-            readout_repetitions=readout_repetitions,
-            qubits=readout_qubits,
+        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
+            sampler, input_circuits, readout_benchmarking_params, rng
         )
         for measurement in measurements:
             assert isinstance(measurement, ResultDict)
@@ -226,14 +225,7 @@ def test_circuits_with_readout_benchmarking_errors_with_noise_and_input_qubits(m
     elif mode == "sweep":
         input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 6)
         sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            input_circuits,
-            sweep_params,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=num_random_bitstrings,
-            readout_repetitions=readout_repetitions,
-            qubits=readout_qubits,
+            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
         )
         for measurement_group in sweep_measurements:  # Treat as a list of lists
             for single_sweep_measurement in measurement_group:
@@ -263,18 +255,20 @@ def test_circuits_with_readout_benchmarking_errors_with_noise_and_lists_input_qu
     circuit_repetitions = 1
     rng = np.random.default_rng()
     readout_repetitions = 1000
+    num_random_bitstrings = 100
+
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=circuit_repetitions,
+        num_random_bitstrings=num_random_bitstrings,
+        readout_repetitions=readout_repetitions,
+        qubits=readout_qubits,
+    )
 
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits_1, 6) + _create_test_circuits(qubits_2, 4)
 
-        measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
-            input_circuits,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=100,
-            readout_repetitions=readout_repetitions,
-            qubits=readout_qubits,
+        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
+            sampler, input_circuits, readout_benchmarking_params, rng
         )
 
         for measurement in measurements:
@@ -287,14 +281,7 @@ def test_circuits_with_readout_benchmarking_errors_with_noise_and_lists_input_qu
         sweep_params += additional_sweep_params
 
         sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            input_circuits,
-            sweep_params,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=100,
-            readout_repetitions=readout_repetitions,
-            qubits=readout_qubits,
+            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
         )
 
         for measurement_group in sweep_measurements:  # Treat as a list of lists
@@ -325,18 +312,20 @@ def test_can_handle_zero_random_bitstring(mode: str):
     circuit_repetitions = 1
     rng = np.random.default_rng()
     readout_repetitions = 1000
+    num_random_bitstrings = 0
+
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=circuit_repetitions,
+        num_random_bitstrings=num_random_bitstrings,
+        readout_repetitions=readout_repetitions,
+        qubits=readout_qubits,
+    )
 
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits_1, 6) + _create_test_circuits(qubits_2, 4)
 
-        measurements, readout_calibration_results = run_shuffled_with_readout_benchmarking(
-            input_circuits,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=0,
-            readout_repetitions=readout_repetitions,
-            qubits=readout_qubits,
+        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
+            sampler, input_circuits, readout_benchmarking_params, rng
         )
 
         for measurement in measurements:
@@ -349,14 +338,7 @@ def test_can_handle_zero_random_bitstring(mode: str):
         sweep_params += additional_sweep_params
 
         sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            input_circuits,
-            sweep_params,
-            sampler,
-            circuit_repetitions,
-            rng,
-            num_random_bitstrings=0,
-            readout_repetitions=readout_repetitions,
-            qubits=readout_qubits,
+            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
         )
 
         for sweep_measurement in sweep_measurements:
@@ -369,28 +351,30 @@ def test_can_handle_zero_random_bitstring(mode: str):
 
 def test_empty_input_circuits():
     """Test that the input circuits are empty."""
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100
+    )
     with pytest.raises(ValueError, match="Input circuits must not be empty."):
-        run_shuffled_with_readout_benchmarking(
-            [],
+        run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
-            circuit_repetitions=10,
+            [],
+            readout_benchmarking_params,
             rng_or_seed=np.random.default_rng(456),
-            num_random_bitstrings=5,
-            readout_repetitions=100,
         )
 
 
 def test_non_circuit_input():
     """Test that the input circuits are not of type cirq.Circuit."""
     q = cirq.LineQubit(0)
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    )
     with pytest.raises(ValueError, match="Input circuits must be of type cirq.Circuit."):
-        run_shuffled_with_readout_benchmarking(
-            [q],
+        run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
-            circuit_repetitions=10,
+            [q],
+            readout_benchmarking_params,
             rng_or_seed=np.random.default_rng(456),
-            num_random_bitstrings=5,
-            readout_repetitions=100,
         )
 
 
@@ -398,29 +382,25 @@ def test_no_measurements():
     """Test that the input circuits don't have measurements."""
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q))
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    )
     with pytest.raises(ValueError, match="Input circuits must have measurements."):
-        run_shuffled_with_readout_benchmarking(
-            [circuit],
+        run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
-            circuit_repetitions=10,
+            [circuit],
+            readout_benchmarking_params,
             rng_or_seed=np.random.default_rng(456),
-            num_random_bitstrings=5,
-            readout_repetitions=100,
         )
 
 
 def test_zero_circuit_repetitions():
     """Test that the circuit repetitions are zero."""
     q = cirq.LineQubit(0)
-    circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
+
     with pytest.raises(ValueError, match="Must provide non-zero circuit_repetitions."):
-        run_shuffled_with_readout_benchmarking(
-            [circuit],
-            cirq.ZerosSampler(),
-            circuit_repetitions=0,
-            rng_or_seed=np.random.default_rng(456),
-            num_random_bitstrings=5,
-            readout_repetitions=100,
+        ReadoutBenchmarkingParams(
+            circuit_repetitions=0, num_random_bitstrings=5, readout_repetitions=100, qubits=q
         )
 
 
@@ -428,32 +408,27 @@ def test_mismatch_circuit_repetitions():
     """Test that the number of circuit repetitions don't match the number of input circuits."""
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=[10, 20], num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    )
     with pytest.raises(
         ValueError,
         match="Number of circuit_repetitions must match the number of" + " input circuits.",
     ):
-        run_shuffled_with_readout_benchmarking(
-            [circuit],
+        run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
-            circuit_repetitions=[10, 20],
+            [circuit],
+            readout_benchmarking_params,
             rng_or_seed=np.random.default_rng(456),
-            num_random_bitstrings=5,
-            readout_repetitions=100,
         )
 
 
 def test_zero_num_random_bitstrings():
     """Test that the number of random bitstrings is smaller than zero."""
     q = cirq.LineQubit(0)
-    circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
     with pytest.raises(ValueError, match="Must provide zero or more num_random_bitstrings."):
-        run_shuffled_with_readout_benchmarking(
-            [circuit],
-            cirq.ZerosSampler(),
-            circuit_repetitions=10,
-            rng_or_seed=np.random.default_rng(456),
-            num_random_bitstrings=-1,
-            readout_repetitions=100,
+        ReadoutBenchmarkingParams(
+            circuit_repetitions=10, num_random_bitstrings=-1, readout_repetitions=100, qubits=q
         )
 
 
@@ -464,13 +439,8 @@ def test_zero_readout_repetitions():
     with pytest.raises(
         ValueError, match="Must provide non-zero readout_repetitions for readout" + " calibration."
     ):
-        run_shuffled_with_readout_benchmarking(
-            [circuit],
-            cirq.ZerosSampler(),
-            circuit_repetitions=10,
-            rng_or_seed=np.random.default_rng(456),
-            num_random_bitstrings=5,
-            readout_repetitions=0,
+        ReadoutBenchmarkingParams(
+            circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=0, qubits=q
         )
 
 
@@ -478,14 +448,15 @@ def test_rng_type_mismatch():
     """Test that the rng is not a numpy random generator or a seed."""
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    )
     with pytest.raises(ValueError, match="Must provide a numpy random generator or a seed"):
-        run_shuffled_with_readout_benchmarking(
-            [circuit],
+        run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
-            circuit_repetitions=10,
+            [circuit],
+            readout_benchmarking_params,
             rng_or_seed="not a random generator or seed",
-            num_random_bitstrings=5,
-            readout_repetitions=100,
         )
 
 
@@ -493,13 +464,14 @@ def test_empty_sweep_params():
     """Test that the sweep params are empty."""
     q = cirq.LineQubit(5)
     circuit = cirq.Circuit(cirq.H(q))
+    readout_benchmarking_params = ReadoutBenchmarkingParams(
+        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    )
     with pytest.raises(ValueError, match="Sweep parameters must not be empty."):
         run_sweep_with_readout_benchmarking(
+            cirq.ZerosSampler(),
             [circuit],
             [],
-            cirq.ZerosSampler(),
-            circuit_repetitions=10,
+            readout_benchmarking_params,
             rng_or_seed=np.random.default_rng(456),
-            num_random_bitstrings=5,
-            readout_repetitions=100,
         )

From 9d56e006e1c54136850e99c3aa09bfac8ea95c1a Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Tue, 1 Jul 2025 23:29:44 -0700
Subject: [PATCH 3/8] Fix tests and lint

---
 ...ing_measurement_with_readout_mitigation.py | 19 +++++----
 ...ffle_circuits_with_readout_benchmarking.py | 40 +++++++++----------
 ...circuits_with_readout_benchmarking_test.py | 10 ++---
 3 files changed, 37 insertions(+), 32 deletions(-)

diff --git a/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation.py b/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation.py
index 0928db748f3..820ec6b1ee9 100644
--- a/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation.py
+++ b/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation.py
@@ -24,7 +24,10 @@
 import numpy as np
 
 from cirq import circuits, ops, work
-from cirq.contrib.shuffle_circuits import run_shuffled_with_readout_benchmarking
+from cirq.contrib.shuffle_circuits import run_shuffled_circuits_with_readout_benchmarking
+from cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking import (
+    ReadoutBenchmarkingParams,
+)
 from cirq.experiments.readout_confusion_matrix import TensoredConfusionMatrices
 
 if TYPE_CHECKING:
@@ -473,14 +476,16 @@ def measure_pauli_strings(
         pauli_measurement_circuits.extend(basis_change_circuits)
 
     # Run shuffled benchmarking for readout calibration
-    circuits_results, calibration_results = run_shuffled_with_readout_benchmarking(
-        input_circuits=pauli_measurement_circuits,
+    circuits_results, calibration_results = run_shuffled_circuits_with_readout_benchmarking(
         sampler=sampler,
-        circuit_repetitions=pauli_repetitions,
+        input_circuits=pauli_measurement_circuits,
+        parameters=ReadoutBenchmarkingParams(
+            circuit_repetitions=pauli_repetitions,
+            num_random_bitstrings=num_random_bitstrings,
+            readout_repetitions=readout_repetitions,
+            qubits=[list(qubits) for qubits in qubits_list],
+        ),
         rng_or_seed=rng_or_seed,
-        qubits=[list(qubits) for qubits in qubits_list],
-        num_random_bitstrings=num_random_bitstrings,
-        readout_repetitions=readout_repetitions,
     )
 
     # Process the results to calculate expectation values
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
index 0056dc39082..66e0c5ddf34 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
@@ -17,15 +17,15 @@
 from __future__ import annotations
 
 import time
-from typing import Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+from typing import Optional, Sequence, TYPE_CHECKING
 
 import attrs
 import numpy as np
 import sympy
 
 from cirq import circuits, ops, protocols, study, work
-from cirq.experiments import SingleQubitReadoutCalibrationResult
 from cirq._compat import deprecated
+from cirq.experiments import SingleQubitReadoutCalibrationResult
 
 if TYPE_CHECKING:
     from cirq.study import ResultDict
@@ -45,10 +45,10 @@ class ReadoutBenchmarkingParams:
                 of qubits.
     """
 
-    circuit_repetitions: Union[int, list[int]]
+    circuit_repetitions: int | list[int]
     num_random_bitstrings: int = 100
     readout_repetitions: int = 1000
-    qubits: Optional[Union[Sequence[ops.Qid], Sequence[Sequence[ops.Qid]]]] = None
+    qubits: Optional[Sequence[ops.Qid] | Sequence[Sequence[ops.Qid]]] = None
 
     def __attrs_post_init__(self):
         _validate_benchmarking_setup(
@@ -57,7 +57,7 @@ def __attrs_post_init__(self):
 
 
 def _validate_benchmarking_setup(
-    circuit_repetitions: Union[int, list[int]], num_random_bitstrings: int, readout_repetitions: int
+    circuit_repetitions: int | list[int], num_random_bitstrings: int, readout_repetitions: int
 ):
     # Check circuit_repetitions
     if isinstance(circuit_repetitions, int):
@@ -75,8 +75,8 @@ def _validate_benchmarking_setup(
 
 def _validate_experiment_input(
     input_circuits: Sequence[circuits.Circuit],
-    circuit_repetitions: Union[int, list[int]],
-    rng_or_seed: Union[np.random.Generator, int],
+    circuit_repetitions: int | list[int],
+    rng_or_seed: np.random.Generator | int,
 ):
     if not input_circuits:
         raise ValueError("Input circuits must not be empty.")
@@ -99,8 +99,8 @@ def _validate_experiment_input(
 def _validate_experiment_input_with_sweep(
     input_circuits: Sequence[circuits.Circuit],
     sweep_params: Sequence[study.Sweepable],
-    circuit_repetitions: Union[int, list[int]],
-    rng_or_seed: Union[np.random.Generator, int],
+    circuit_repetitions: int | list[int],
+    rng_or_seed: np.random.Generator | int,
 ):
     """Validates the input for the run_sweep_with_readout_benchmarking function."""
     if not sweep_params:
@@ -164,10 +164,10 @@ def _generate_parameterized_readout_calibration_circuit_with_sweep(
 
 def _generate_all_readout_calibration_circuits(
     num_random_bitstrings: int,
-    qubits_to_measure: List[List[ops.Qid]],
+    qubits_to_measure: list[list[ops.Qid]],
     is_sweep: bool,
     rng: np.random.Generator,
-) -> Tuple[List[circuits.Circuit], List[np.ndarray], List[study.Sweepable]]:
+) -> tuple[list[circuits.Circuit], list[np.ndarray], list[study.Sweepable]]:
     """Generates all readout calibration circuits and random bitstrings."""
     all_readout_calibration_circuits: list[circuits.Circuit] = []
     all_random_bitstrings: list[np.ndarray] = []
@@ -199,11 +199,11 @@ def _generate_all_readout_calibration_circuits(
 
 def _determine_qubits_to_measure(
     input_circuits: Sequence[circuits.Circuit],
-    qubits: Optional[Union[Sequence[ops.Qid], Sequence[Sequence[ops.Qid]]]],
-) -> List[List[ops.Qid]]:
+    qubits: Optional[Sequence[ops.Qid] | Sequence[Sequence[ops.Qid]]],
+) -> list[list[ops.Qid]]:
     """Determine the qubits to measure based on the input circuits and provided qubits."""
     # If input qubits is None, extract qubits from input circuits
-    qubits_to_measure: List[List[ops.Qid]] = []
+    qubits_to_measure: list[list[ops.Qid]] = []
     if qubits is None:
         qubits_to_measure = [
             sorted(set(q for circuit in input_circuits for q in circuit.all_qubits()))
@@ -229,7 +229,7 @@ def _shuffle_circuits(
 
 
 def _analyze_readout_results(
-    unshuffled_readout_measurements: Union[Sequence[ResultDict], Sequence[study.Result]],
+    unshuffled_readout_measurements: list[ResultDict] | list[study.Result],
     random_bitstrings: np.ndarray,
     readout_repetitions: int,
     qubits: list[ops.Qid],
@@ -296,8 +296,8 @@ def run_shuffled_with_readout_benchmarking(
     rng_or_seed: np.random.Generator | int,
     num_random_bitstrings: int = 100,
     readout_repetitions: int = 1000,
-    qubits: Optional[Union[Sequence[ops.Qid], Sequence[Sequence[ops.Qid]]]] = None,
-) -> tuple[Sequence[ResultDict], Dict[Tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]:
+    qubits: Optional[Sequence[ops.Qid] | Sequence[Sequence[ops.Qid]]] = None,
+) -> tuple[Sequence[ResultDict], dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]:
     """Run the circuits in a shuffled order with readout error benchmarking.
 
     Args:
@@ -379,7 +379,7 @@ def run_shuffled_circuits_with_readout_benchmarking(
     input_circuits: list[circuits.Circuit],
     parameters: ReadoutBenchmarkingParams,
     rng_or_seed: np.random.Generator | int,
-) -> tuple[Sequence[ResultDict], Dict[Tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]:
+) -> tuple[Sequence[ResultDict], dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]:
     """Run the circuits in a shuffled order with readout error benchmarking.
 
     Args:
@@ -459,9 +459,9 @@ def run_sweep_with_readout_benchmarking(
     input_circuits: list[circuits.Circuit],
     sweep_params: Sequence[study.Sweepable],
     parameters: ReadoutBenchmarkingParams,
-    rng_or_seed: Union[np.random.Generator, int],
+    rng_or_seed: np.random.Generator | int,
 ) -> tuple[
-    Sequence[Sequence[study.Result]], Dict[Tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]
+    Sequence[Sequence[study.Result]], dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]
 ]:
     """Run the sweep circuits with readout error benchmarking (no shuffling).
     Args:
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
index 958aa9062b4..ccdb9ecfcb3 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
@@ -17,16 +17,15 @@
 import itertools
 from typing import Sequence
 
-import cirq.contrib.shuffle_circuits
 import numpy as np
 import pytest
 import sympy
 
 import cirq
+import cirq.contrib.shuffle_circuits
 from cirq.contrib.shuffle_circuits import (
     run_shuffled_circuits_with_readout_benchmarking,
     run_sweep_with_readout_benchmarking,
-    run_shuffled_with_readout_benchmarking,
 )
 from cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking import (
     ReadoutBenchmarkingParams,
@@ -196,7 +195,8 @@ def test_circuits_with_readout_benchmarking_errors_with_noise(mode: str):
 
 @pytest.mark.parametrize("mode", ["shuffled", "sweep"])
 def test_circuits_with_readout_benchmarking_errors_with_noise_and_input_qubits(mode: str):
-    """Test shuffled/sweep circuits with readout benchmarking with noise from sampler and input qubits."""
+    """Test shuffled/sweep circuits with readout benchmarking with
+    noise from sampler and input qubits."""
     qubits = cirq.LineQubit.range(6)
     readout_qubits = qubits[:4]
 
@@ -246,7 +246,8 @@ def test_circuits_with_readout_benchmarking_errors_with_noise_and_input_qubits(m
 
 @pytest.mark.parametrize("mode", ["shuffled", "sweep"])
 def test_circuits_with_readout_benchmarking_errors_with_noise_and_lists_input_qubits(mode: str):
-    """Test shuffled/sweep circuits with readout benchmarking with noise from sampler and input qubits."""
+    """Test shuffled/sweep circuits with readout benchmarking with noise
+    from sampler and input qubits."""
     qubits_1 = cirq.LineQubit.range(3)
     qubits_2 = cirq.LineQubit.range(4)
     readout_qubits = [qubits_1, qubits_2]
@@ -435,7 +436,6 @@ def test_zero_num_random_bitstrings():
 def test_zero_readout_repetitions():
     """Test that the readout repetitions is zero."""
     q = cirq.LineQubit(0)
-    circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
     with pytest.raises(
         ValueError, match="Must provide non-zero readout_repetitions for readout" + " calibration."
     ):

From 236ce12129f4bea3c308017ca6374af7c92355b7 Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Thu, 24 Jul 2025 18:33:32 -0700
Subject: [PATCH 4/8] Resolve comments

---
 ...ing_measurement_with_readout_mitigation.py |  29 ++-
 ...ffle_circuits_with_readout_benchmarking.py |  18 +-
 ...circuits_with_readout_benchmarking_test.py | 182 +++++++++---------
 3 files changed, 113 insertions(+), 116 deletions(-)

diff --git a/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation.py b/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation.py
index 820ec6b1ee9..cad69a3d4b2 100644
--- a/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation.py
+++ b/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation.py
@@ -23,11 +23,8 @@
 import attrs
 import numpy as np
 
-from cirq import circuits, ops, work
-from cirq.contrib.shuffle_circuits import run_shuffled_circuits_with_readout_benchmarking
-from cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking import (
-    ReadoutBenchmarkingParams,
-)
+import cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking as sc_readout
+from cirq import circuits, ops, study, work
 from cirq.experiments.readout_confusion_matrix import TensoredConfusionMatrices
 
 if TYPE_CHECKING:
@@ -291,7 +288,7 @@ def _build_many_one_qubits_empty_confusion_matrix(qubits_length: int) -> list[np
 def _process_pauli_measurement_results(
     qubits: list[ops.Qid],
     pauli_string_groups: list[list[ops.PauliString]],
-    circuit_results: list[ResultDict],
+    circuit_results: list[ResultDict] | Sequence[study.Result],
     calibration_results: dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult],
     pauli_repetitions: int,
     timestamp: float,
@@ -476,16 +473,18 @@ def measure_pauli_strings(
         pauli_measurement_circuits.extend(basis_change_circuits)
 
     # Run shuffled benchmarking for readout calibration
-    circuits_results, calibration_results = run_shuffled_circuits_with_readout_benchmarking(
-        sampler=sampler,
-        input_circuits=pauli_measurement_circuits,
-        parameters=ReadoutBenchmarkingParams(
-            circuit_repetitions=pauli_repetitions,
-            num_random_bitstrings=num_random_bitstrings,
-            readout_repetitions=readout_repetitions,
+    circuits_results, calibration_results = (
+        sc_readout.run_shuffled_circuits_with_readout_benchmarking(
+            sampler=sampler,
+            input_circuits=pauli_measurement_circuits,
+            parameters=sc_readout.ReadoutBenchmarkingParams(
+                circuit_repetitions=pauli_repetitions,
+                num_random_bitstrings=num_random_bitstrings,
+                readout_repetitions=readout_repetitions,
+            ),
+            rng_or_seed=rng_or_seed,
             qubits=[list(qubits) for qubits in qubits_list],
-        ),
-        rng_or_seed=rng_or_seed,
+        )
     )
 
     # Process the results to calculate expectation values
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
index 66e0c5ddf34..e7227800b2a 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
@@ -40,15 +40,11 @@ class ReadoutBenchmarkingParams:
         num_random_bitstrings: The number of random bitstrings for measuring readout.
             If set to 0, no readout calibration circuits are generated.
         readout_repetitions: The number of repetitions for each readout bitstring.
-        qubits: The qubits to benchmark readout errors. If None, all qubits in the
-                input_circuits are used. Can be a list of qubits or a list of tuples
-                of qubits.
     """
 
     circuit_repetitions: int | list[int]
     num_random_bitstrings: int = 100
     readout_repetitions: int = 1000
-    qubits: Optional[Sequence[ops.Qid] | Sequence[Sequence[ops.Qid]]] = None
 
     def __attrs_post_init__(self):
         _validate_benchmarking_setup(
@@ -229,7 +225,7 @@ def _shuffle_circuits(
 
 
 def _analyze_readout_results(
-    unshuffled_readout_measurements: list[ResultDict] | list[study.Result],
+    unshuffled_readout_measurements: Sequence[ResultDict] | Sequence[study.Result],
     random_bitstrings: np.ndarray,
     readout_repetitions: int,
     qubits: list[ops.Qid],
@@ -379,6 +375,7 @@ def run_shuffled_circuits_with_readout_benchmarking(
     input_circuits: list[circuits.Circuit],
     parameters: ReadoutBenchmarkingParams,
     rng_or_seed: np.random.Generator | int,
+    qubits: Optional[Sequence[ops.Qid] | Sequence[Sequence[ops.Qid]]] = None,
 ) -> tuple[Sequence[ResultDict], dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]:
     """Run the circuits in a shuffled order with readout error benchmarking.
 
@@ -386,6 +383,9 @@ def run_shuffled_circuits_with_readout_benchmarking(
         sampler: The sampler to use.
         input_circuits: The circuits to run.
         parameters: The readout benchmarking parameters.
+        qubits: The qubits to benchmark readout errors. If None, all qubits in the
+                input_circuits are used. Can be a list of qubits or a list of tuples
+                of qubits.
         rng_or_seed: A random number generator used to generate readout circuits.
                      Or an integer seed.
 
@@ -398,7 +398,7 @@ def run_shuffled_circuits_with_readout_benchmarking(
 
     _validate_experiment_input(input_circuits, parameters.circuit_repetitions, rng_or_seed)
 
-    qubits_to_measure = _determine_qubits_to_measure(input_circuits, parameters.qubits)
+    qubits_to_measure = _determine_qubits_to_measure(input_circuits, qubits)
 
     # Generate the readout calibration circuits if num_random_bitstrings>0
     # Else all_readout_calibration_circuits and all_random_bitstrings are empty
@@ -460,6 +460,7 @@ def run_sweep_with_readout_benchmarking(
     sweep_params: Sequence[study.Sweepable],
     parameters: ReadoutBenchmarkingParams,
     rng_or_seed: np.random.Generator | int,
+    qubits: Optional[Sequence[ops.Qid] | Sequence[Sequence[ops.Qid]]] = None,
 ) -> tuple[
     Sequence[Sequence[study.Result]], dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]
 ]:
@@ -471,6 +472,9 @@ def run_sweep_with_readout_benchmarking(
         parameters: The readout benchmarking parameters.
         rng_or_seed: A random number generator used to generate readout circuits.
                      Or an integer seed.
+        qubits: The qubits to benchmark readout errors. If None, all qubits in the
+        input_circuits are used. Can be a list of qubits or a list of tuples
+        of qubits.
 
     Returns:
         A tuple containing:
@@ -482,7 +486,7 @@ def run_sweep_with_readout_benchmarking(
         input_circuits, sweep_params, parameters.circuit_repetitions, rng_or_seed
     )
 
-    qubits_to_measure = _determine_qubits_to_measure(input_circuits, parameters.qubits)
+    qubits_to_measure = _determine_qubits_to_measure(input_circuits, qubits)
 
     # Generate the readout calibration circuits (parameterized circuits) and sweep params
     # if num_random_bitstrings>0
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
index ccdb9ecfcb3..6c5af58768d 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
@@ -22,14 +22,7 @@
 import sympy
 
 import cirq
-import cirq.contrib.shuffle_circuits
-from cirq.contrib.shuffle_circuits import (
-    run_shuffled_circuits_with_readout_benchmarking,
-    run_sweep_with_readout_benchmarking,
-)
-from cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking import (
-    ReadoutBenchmarkingParams,
-)
+import cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking as sc_readout
 from cirq.experiments import (
     random_quantum_circuit_generation as rqcg,
     SingleQubitReadoutCalibrationResult,
@@ -92,6 +85,44 @@ def _create_test_circuits_with_sweep(
     return input_circuits, sweep_params
 
 
+def _circuits_with_readout_benchmarking_errors_shuffled(
+    sampler: cirq.Sampler,
+    input_circuits: list[cirq.Circuit],
+    qubits: Sequence[cirq.Qid] | Sequence[Sequence[cirq.Qid]],
+    parameters: sc_readout.ReadoutBenchmarkingParams,
+    rng_or_seed: np.random.Generator | int,
+):
+    measurements, readout_calibration_results = (
+        sc_readout.run_shuffled_circuits_with_readout_benchmarking(
+            sampler, input_circuits, parameters, rng_or_seed, qubits
+        )
+    )
+
+    for measurement in measurements:
+        assert isinstance(measurement, ResultDict)
+    return readout_calibration_results
+
+
+def _circuits_with_readout_benchmarking_errors_sweep(
+    sampler: cirq.Sampler,
+    input_circuits: list[cirq.Circuit],
+    qubits: Sequence[cirq.Qid] | Sequence[Sequence[cirq.Qid]],
+    sweep_params: list[cirq.ParamResolver],
+    parameters: sc_readout.ReadoutBenchmarkingParams,
+    rng_or_seed: np.random.Generator | int,
+):
+    sweep_measurements, readout_calibration_results = (
+        sc_readout.run_sweep_with_readout_benchmarking(
+            sampler, input_circuits, sweep_params, parameters, rng_or_seed, qubits
+        )
+    )
+
+    for measurement_group in sweep_measurements:
+        for single_sweep_measurement in measurement_group:
+            assert isinstance(single_sweep_measurement, ResultDict)
+    return readout_calibration_results
+
+
 @pytest.mark.parametrize("mode", ["shuffled", "sweep"])
 def test_circuits_with_readout_benchmarking_errors_no_noise(mode: str):
     """Test shuffled/sweep circuits with readout benchmarking with no noise from sampler."""
@@ -102,36 +133,27 @@ def test_circuits_with_readout_benchmarking_errors_no_noise(mode: str):
     num_random_bitstrings = 100
     readout_repetitions = 1000
 
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
         circuit_repetitions=circuit_repetitions,
         num_random_bitstrings=num_random_bitstrings,
         readout_repetitions=readout_repetitions,
-        qubits=qubits,
     )
     # allow passing a seed
     rng = 123
 
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits, 3)
-
-        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
-            sampler, input_circuits, readout_benchmarking_params, rng
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_shuffled(
+            sampler, input_circuits, qubits, readout_benchmarking_params, rng
         )
 
-        for measurement in measurements:
-            assert isinstance(measurement, ResultDict)
-
     elif mode == "sweep":
         input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 3)
 
-        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_sweep(
+            sampler, input_circuits, qubits, sweep_params, readout_benchmarking_params, rng
         )
 
-        for measurement_group in sweep_measurements:
-            for single_sweep_measurement in measurement_group:
-                assert isinstance(single_sweep_measurement, ResultDict)
-
     for qlist, readout_calibration_result in readout_calibration_results.items():
         assert isinstance(qlist, tuple)
         assert all(isinstance(q, cirq.Qid) for q in qlist)
@@ -153,33 +175,26 @@ def test_circuits_with_readout_benchmarking_errors_with_noise(mode: str):
     readout_repetitions = 1000
     num_random_bitstrings = 100
 
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
         circuit_repetitions=circuit_repetitions,
         num_random_bitstrings=num_random_bitstrings,
         readout_repetitions=readout_repetitions,
-        qubits=qubits,
     )
 
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits, 6)
 
-        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
-            sampler, input_circuits, readout_benchmarking_params, rng
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_shuffled(
+            sampler, input_circuits, qubits, readout_benchmarking_params, rng
         )
 
-        for measurement in measurements:
-            assert isinstance(measurement, ResultDict)
     elif mode == "sweep":
         input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 6)
 
-        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_sweep(
+            sampler, input_circuits, qubits, sweep_params, readout_benchmarking_params, rng
         )
 
-        for measurement_group in sweep_measurements:  # Treat as a list of lists
-            for single_sweep_measurement in measurement_group:
-                assert isinstance(single_sweep_measurement, ResultDict)
-
     for qlist, readout_calibration_result in readout_calibration_results.items():
         assert isinstance(qlist, tuple)
         assert all(isinstance(q, cirq.Qid) for q in qlist)
@@ -206,30 +221,25 @@ def test_circuits_with_readout_benchmarking_errors_with_noise_and_input_qubits(m
     readout_repetitions = 1000
     num_random_bitstrings = 100
 
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
         circuit_repetitions=circuit_repetitions,
         num_random_bitstrings=num_random_bitstrings,
         readout_repetitions=readout_repetitions,
-        qubits=readout_qubits,
     )
 
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits, 6)
 
-        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
-            sampler, input_circuits, readout_benchmarking_params, rng
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_shuffled(
+            sampler, input_circuits, readout_qubits, readout_benchmarking_params, rng
         )
-        for measurement in measurements:
-            assert isinstance(measurement, ResultDict)
 
     elif mode == "sweep":
         input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 6)
-        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
+
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_sweep(
+            sampler, input_circuits, readout_qubits, sweep_params, readout_benchmarking_params, rng
         )
-        for measurement_group in sweep_measurements:  # Treat as a list of lists
-            for single_sweep_measurement in measurement_group:
-                assert isinstance(single_sweep_measurement, ResultDict)
 
     for qlist, readout_calibration_result in readout_calibration_results.items():
         assert isinstance(qlist, tuple)
@@ -258,37 +268,29 @@ def test_circuits_with_readout_benchmarking_errors_with_noise_and_lists_input_qu
     readout_repetitions = 1000
     num_random_bitstrings = 100
 
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
         circuit_repetitions=circuit_repetitions,
         num_random_bitstrings=num_random_bitstrings,
         readout_repetitions=readout_repetitions,
-        qubits=readout_qubits,
     )
 
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits_1, 6) + _create_test_circuits(qubits_2, 4)
 
-        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
-            sampler, input_circuits, readout_benchmarking_params, rng
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_shuffled(
+            sampler, input_circuits, readout_qubits, readout_benchmarking_params, rng
         )
 
-        for measurement in measurements:
-            assert isinstance(measurement, ResultDict)
-
     elif mode == "sweep":
         input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits_1, 6)
         additional_circuits, additional_sweep_params = _create_test_circuits_with_sweep(qubits_2, 4)
         input_circuits += additional_circuits
         sweep_params += additional_sweep_params
 
-        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_sweep(
+            sampler, input_circuits, readout_qubits, sweep_params, readout_benchmarking_params, rng
         )
 
-        for measurement_group in sweep_measurements:  # Treat as a list of lists
-            for single_sweep_measurement in measurement_group:
-                assert isinstance(single_sweep_measurement, ResultDict)
-
     for qlist, readout_calibration_result in readout_calibration_results.items():
         assert isinstance(qlist, tuple)
         assert all(isinstance(q, cirq.Qid) for q in qlist)
@@ -315,48 +317,40 @@ def test_can_handle_zero_random_bitstring(mode: str):
     readout_repetitions = 1000
     num_random_bitstrings = 0
 
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
         circuit_repetitions=circuit_repetitions,
         num_random_bitstrings=num_random_bitstrings,
         readout_repetitions=readout_repetitions,
-        qubits=readout_qubits,
     )
 
     if mode == "shuffled":
         input_circuits = _create_test_circuits(qubits_1, 6) + _create_test_circuits(qubits_2, 4)
 
-        measurements, readout_calibration_results = run_shuffled_circuits_with_readout_benchmarking(
-            sampler, input_circuits, readout_benchmarking_params, rng
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_shuffled(
+            sampler, input_circuits, readout_qubits, readout_benchmarking_params, rng
         )
 
-        for measurement in measurements:
-            assert isinstance(measurement, ResultDict)
-
     elif mode == "sweep":
         input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits_1, 6)
         additional_circuits, additional_sweep_params = _create_test_circuits_with_sweep(qubits_2, 4)
         input_circuits += additional_circuits
         sweep_params += additional_sweep_params
 
-        sweep_measurements, readout_calibration_results = run_sweep_with_readout_benchmarking(
-            sampler, input_circuits, sweep_params, readout_benchmarking_params, rng
+        readout_calibration_results = _circuits_with_readout_benchmarking_errors_sweep(
+            sampler, input_circuits, readout_qubits, sweep_params, readout_benchmarking_params, rng
         )
 
-        for sweep_measurement in sweep_measurements:
-            for single_sweep_measurement in sweep_measurement:
-                assert isinstance(single_sweep_measurement, ResultDict)
-
     # Check that the readout_calibration_results is empty
     assert len(readout_calibration_results.items()) == 0
 
 
 def test_empty_input_circuits():
     """Test that the input circuits are empty."""
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
         circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100
     )
     with pytest.raises(ValueError, match="Input circuits must not be empty."):
-        run_shuffled_circuits_with_readout_benchmarking(
+        sc_readout.run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
             [],
             readout_benchmarking_params,
@@ -367,11 +361,11 @@ def test_empty_input_circuits():
 def test_non_circuit_input():
     """Test that the input circuits are not of type cirq.Circuit."""
     q = cirq.LineQubit(0)
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
-        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
+        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100
     )
     with pytest.raises(ValueError, match="Input circuits must be of type cirq.Circuit."):
-        run_shuffled_circuits_with_readout_benchmarking(
+        sc_readout.run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
             [q],
             readout_benchmarking_params,
@@ -383,11 +377,11 @@ def test_no_measurements():
     """Test that the input circuits don't have measurements."""
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q))
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
-        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
+        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100
     )
     with pytest.raises(ValueError, match="Input circuits must have measurements."):
-        run_shuffled_circuits_with_readout_benchmarking(
+        sc_readout.run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
             [circuit],
             readout_benchmarking_params,
@@ -400,8 +394,8 @@ def test_zero_circuit_repetitions():
     q = cirq.LineQubit(0)
 
     with pytest.raises(ValueError, match="Must provide non-zero circuit_repetitions."):
-        ReadoutBenchmarkingParams(
-            circuit_repetitions=0, num_random_bitstrings=5, readout_repetitions=100, qubits=q
+        sc_readout.ReadoutBenchmarkingParams(
+            circuit_repetitions=0, num_random_bitstrings=5, readout_repetitions=100
         )
 
 
@@ -409,14 +403,14 @@ def test_mismatch_circuit_repetitions():
     """Test that the number of circuit repetitions don't match the number of input circuits."""
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
-        circuit_repetitions=[10, 20], num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
+        circuit_repetitions=[10, 20], num_random_bitstrings=5, readout_repetitions=100
     )
     with pytest.raises(
         ValueError,
         match="Number of circuit_repetitions must match the number of" + " input circuits.",
     ):
-        run_shuffled_circuits_with_readout_benchmarking(
+        sc_readout.run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
             [circuit],
             readout_benchmarking_params,
@@ -428,8 +422,8 @@ def test_zero_num_random_bitstrings():
     """Test that the number of random bitstrings is smaller than zero."""
     q = cirq.LineQubit(0)
     with pytest.raises(ValueError, match="Must provide zero or more num_random_bitstrings."):
-        ReadoutBenchmarkingParams(
-            circuit_repetitions=10, num_random_bitstrings=-1, readout_repetitions=100, qubits=q
+        sc_readout.ReadoutBenchmarkingParams(
+            circuit_repetitions=10, num_random_bitstrings=-1, readout_repetitions=100
         )
 
 
@@ -439,8 +433,8 @@ def test_zero_readout_repetitions():
     with pytest.raises(
         ValueError, match="Must provide non-zero readout_repetitions for readout" + " calibration."
     ):
-        ReadoutBenchmarkingParams(
-            circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=0, qubits=q
+        sc_readout.ReadoutBenchmarkingParams(
+            circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=0
         )
 
 
@@ -448,11 +442,11 @@ def test_rng_type_mismatch():
     """Test that the rng is not a numpy random generator or a seed."""
     q = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
-        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
+        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100
     )
     with pytest.raises(ValueError, match="Must provide a numpy random generator or a seed"):
-        run_shuffled_circuits_with_readout_benchmarking(
+        sc_readout.run_shuffled_circuits_with_readout_benchmarking(
             cirq.ZerosSampler(),
             [circuit],
             readout_benchmarking_params,
@@ -464,11 +458,11 @@ def test_empty_sweep_params():
     """Test that the sweep params are empty."""
     q = cirq.LineQubit(5)
     circuit = cirq.Circuit(cirq.H(q))
-    readout_benchmarking_params = ReadoutBenchmarkingParams(
-        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100, qubits=q
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
+        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100
     )
     with pytest.raises(ValueError, match="Sweep parameters must not be empty."):
-        run_sweep_with_readout_benchmarking(
+        sc_readout.run_sweep_with_readout_benchmarking(
             cirq.ZerosSampler(),
             [circuit],
             [],

From 1b8b10bf8c6d3775f33772eee1cb607be7d07223 Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Thu, 24 Jul 2025 19:29:57 -0700
Subject: [PATCH 5/8] Fix lint and type checks

---
 ...easurement_with_readout_mitigation_test.py |  2 +-
 .../cirq/contrib/shuffle_circuits/__init__.py |  2 +-
 ...circuits_with_readout_benchmarking_test.py | 93 ++++++++++++++++++-
 3 files changed, 91 insertions(+), 6 deletions(-)

diff --git a/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation_test.py b/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation_test.py
index 1992ead250d..a2000d4e0b6 100644
--- a/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation_test.py
+++ b/cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation_test.py
@@ -899,7 +899,7 @@ def test_process_pauli_measurement_results_raises_error_on_missing_calibration()
         _process_pauli_measurement_results(
             qubits,
             [pauli_strings],
-            circuit_results[0],  # type: ignore[arg-type]
+            circuit_results[0],
             empty_calibration_result_dict,  # type: ignore[arg-type]
             1000,
             1.0,
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/__init__.py b/cirq-core/cirq/contrib/shuffle_circuits/__init__.py
index a2b90c4f96b..2c1a29702df 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/__init__.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/__init__.py
@@ -15,6 +15,6 @@
 
 from cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking import (
     run_shuffled_with_readout_benchmarking as run_shuffled_with_readout_benchmarking,
-    run_shuffled_circuits_with_readout_benchmarking as run_shuffled_circuits_with_readout_benchmarking,
+    run_shuffled_circuits_with_readout_benchmarking as run_shuffled_circuits_with_readout_benchmarking,  # noqa: E501
     run_sweep_with_readout_benchmarking as run_sweep_with_readout_benchmarking,
 )
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
index 6c5af58768d..fbb98ac895f 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
@@ -344,6 +344,95 @@ def test_can_handle_zero_random_bitstring(mode: str):
     assert len(readout_calibration_results.items()) == 0
 
 
+@pytest.mark.parametrize("mode", ["shuffled", "sweep"])
+def test_circuits_with_readout_benchmarking_no_qubits_arg(mode: str):
+    """Test benchmarking when the `qubits` argument is not provided."""
+    qubits = cirq.LineQubit.range(3)
+    sampler = NoisySingleQubitReadoutSampler(p0=0.1, p1=0.2, seed=1234)
+    circuit_repetitions = 1
+    rng = np.random.default_rng()
+    readout_repetitions = 1000
+    num_random_bitstrings = 100
+
+    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
+        circuit_repetitions=circuit_repetitions,
+        num_random_bitstrings=num_random_bitstrings,
+        readout_repetitions=readout_repetitions,
+    )
+
+    if mode == "shuffled":
+        input_circuits = _create_test_circuits(qubits, 3)
+        measurements, readout_calibration_results = (
+            sc_readout.run_shuffled_circuits_with_readout_benchmarking(
+                sampler, input_circuits, readout_benchmarking_params, rng, qubits=None
+            )
+        )
+        assert len(measurements) == len(input_circuits)
+    else:  # mode == "sweep"
+        input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 3)
+        sweep_measurements, readout_calibration_results = (
+            sc_readout.run_sweep_with_readout_benchmarking(
+                sampler, input_circuits, sweep_params, readout_benchmarking_params, rng, qubits=None
+            )
+        )
+        assert len(sweep_measurements) == len(input_circuits)
+
+    # When qubits is None, all qubits from input circuits are benchmarked as one group.
+    assert len(readout_calibration_results) == 1
+    qlist, result = list(readout_calibration_results.items())[0]
+    assert isinstance(qlist, tuple)
+    assert set(qlist) == set(qubits)
+    assert isinstance(result, SingleQubitReadoutCalibrationResult)
+    for error in result.zero_state_errors.values():
+        assert 0.08 < error < 0.12
+    for error in result.one_state_errors.values():
+        assert 0.18 < error < 0.22
+    assert result.repetitions == readout_repetitions
+
+
+def test_deprecated_run_shuffled_with_readout_benchmarking():
+    """Test that the deprecated function works correctly and is covered."""
+    qubits = cirq.LineQubit.range(3)
+    input_circuits = _create_test_circuits(qubits, 3)
+    sampler = NoisySingleQubitReadoutSampler(p0=0.1, p1=0.2, seed=1234)
+    circuit_repetitions = 1
+    readout_repetitions = 1000
+    num_random_bitstrings = 100
+
+    # Test with an integer seed.
+    with cirq.testing.assert_deprecated(deadline='v2.0', count=1):
+        measurements_seed, results_seed = sc_readout.run_shuffled_with_readout_benchmarking(
+            input_circuits=input_circuits,
+            sampler=sampler,
+            circuit_repetitions=circuit_repetitions,
+            rng_or_seed=123,
+            num_random_bitstrings=num_random_bitstrings,
+            readout_repetitions=readout_repetitions,
+            qubits=qubits,
+        )
+    assert len(measurements_seed) == len(input_circuits)
+    qlist, result = list(results_seed.items())[0]
+    assert tuple(qubits) == qlist
+    for error in result.zero_state_errors.values():
+        assert 0.08 < error < 0.12
+    for error in result.one_state_errors.values():
+        assert 0.18 < error < 0.22
+
+    # Test with qubits=None to cover the auto-detection branch.
+    with cirq.testing.assert_deprecated(deadline='v2.0', count=1):
+        _, results_none = sc_readout.run_shuffled_with_readout_benchmarking(
+            input_circuits=input_circuits,
+            sampler=sampler,
+            circuit_repetitions=circuit_repetitions,
+            rng_or_seed=123,
+            num_random_bitstrings=num_random_bitstrings,
+            readout_repetitions=readout_repetitions,
+            qubits=None,
+        )
+    qlist_none, _ = list(results_none.items())[0]
+    assert set(qlist_none) == set(qubits)
+
+
 def test_empty_input_circuits():
     """Test that the input circuits are empty."""
     readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
@@ -391,8 +480,6 @@ def test_no_measurements():
 
 def test_zero_circuit_repetitions():
     """Test that the circuit repetitions are zero."""
-    q = cirq.LineQubit(0)
-
     with pytest.raises(ValueError, match="Must provide non-zero circuit_repetitions."):
         sc_readout.ReadoutBenchmarkingParams(
             circuit_repetitions=0, num_random_bitstrings=5, readout_repetitions=100
@@ -420,7 +507,6 @@ def test_mismatch_circuit_repetitions():
 
 def test_zero_num_random_bitstrings():
     """Test that the number of random bitstrings is smaller than zero."""
-    q = cirq.LineQubit(0)
     with pytest.raises(ValueError, match="Must provide zero or more num_random_bitstrings."):
         sc_readout.ReadoutBenchmarkingParams(
             circuit_repetitions=10, num_random_bitstrings=-1, readout_repetitions=100
@@ -429,7 +515,6 @@ def test_zero_num_random_bitstrings():
 
 def test_zero_readout_repetitions():
     """Test that the readout repetitions is zero."""
-    q = cirq.LineQubit(0)
     with pytest.raises(
         ValueError, match="Must provide non-zero readout_repetitions for readout" + " calibration."
     ):

From 17d46e3eb1a2bddf2fa7bd616c34f05da9763671 Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Tue, 29 Jul 2025 13:36:09 -0700
Subject: [PATCH 6/8] Address Nour's comments

---
 ...ffle_circuits_with_readout_benchmarking.py | 66 ++++++++++---------
 ...circuits_with_readout_benchmarking_test.py | 11 ++--
 2 files changed, 41 insertions(+), 36 deletions(-)

diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
index e7227800b2a..66d6b0c4245 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
@@ -47,32 +47,24 @@ class ReadoutBenchmarkingParams:
     readout_repetitions: int = 1000
 
     def __attrs_post_init__(self):
-        _validate_benchmarking_setup(
-            self.circuit_repetitions, self.num_random_bitstrings, self.readout_repetitions
-        )
-
+        # Check circuit_repetitions
+        if isinstance(self.circuit_repetitions, int):
+            if self.circuit_repetitions <= 0:
+                raise ValueError("Must provide non-zero circuit_repetitions.")
 
-def _validate_benchmarking_setup(
-    circuit_repetitions: int | list[int], num_random_bitstrings: int, readout_repetitions: int
-):
-    # Check circuit_repetitions
-    if isinstance(circuit_repetitions, int):
-        if circuit_repetitions <= 0:
-            raise ValueError("Must provide non-zero circuit_repetitions.")
-
-    # Check num_random_bitstrings is bigger than or equal to 0
-    if num_random_bitstrings < 0:
-        raise ValueError("Must provide zero or more num_random_bitstrings.")
+        # Check num_random_bitstrings is bigger than or equal to 0
+        if self.num_random_bitstrings < 0:
+            raise ValueError("Must provide zero or more num_random_bitstrings.")
 
-    # Check readout_repetitions is bigger than 0
-    if readout_repetitions <= 0:
-        raise ValueError("Must provide non-zero readout_repetitions for readout calibration.")
+        # Check readout_repetitions is bigger than 0
+        if self.readout_repetitions <= 0:
+            raise ValueError("Must provide non-zero readout_repetitions for readout calibration.")
 
 
 def _validate_experiment_input(
     input_circuits: Sequence[circuits.Circuit],
     circuit_repetitions: int | list[int],
-    rng_or_seed: np.random.Generator | int,
+    rng_or_seed: Optional[np.random.Generator | int] = None,
 ):
     if not input_circuits:
         raise ValueError("Input circuits must not be empty.")
@@ -88,7 +80,11 @@ def _validate_experiment_input(
         raise ValueError("Number of circuit_repetitions must match the number of input circuits.")
 
     # Check rng is a numpy random generator
-    if not isinstance(rng_or_seed, np.random.Generator) and not isinstance(rng_or_seed, int):
+    if (
+        rng_or_seed
+        and not isinstance(rng_or_seed, np.random.Generator)
+        and not isinstance(rng_or_seed, int)
+    ):
         raise ValueError("Must provide a numpy random generator or a seed")
 
 
@@ -96,7 +92,7 @@ def _validate_experiment_input_with_sweep(
     input_circuits: Sequence[circuits.Circuit],
     sweep_params: Sequence[study.Sweepable],
     circuit_repetitions: int | list[int],
-    rng_or_seed: np.random.Generator | int,
+    rng_or_seed: Optional[np.random.Generator | int] = None,
 ):
     """Validates the input for the run_sweep_with_readout_benchmarking function."""
     if not sweep_params:
@@ -316,7 +312,18 @@ def run_shuffled_with_readout_benchmarking(
 
     """
 
-    _validate_benchmarking_setup(circuit_repetitions, num_random_bitstrings, readout_repetitions)
+    # Check circuit_repetitions
+    if isinstance(circuit_repetitions, int):
+        if circuit_repetitions <= 0:
+            raise ValueError("Must provide non-zero circuit_repetitions.")
+
+    # Check num_random_bitstrings is bigger than or equal to 0
+    if num_random_bitstrings < 0:
+        raise ValueError("Must provide zero or more num_random_bitstrings.")
+
+    # Check readout_repetitions is bigger than 0
+    if readout_repetitions <= 0:
+        raise ValueError("Must provide non-zero readout_repetitions for readout calibration.")
     _validate_experiment_input(input_circuits, circuit_repetitions, rng_or_seed)
 
     qubits_to_measure = _determine_qubits_to_measure(input_circuits, qubits)
@@ -374,8 +381,8 @@ def run_shuffled_circuits_with_readout_benchmarking(
     sampler: work.Sampler,
     input_circuits: list[circuits.Circuit],
     parameters: ReadoutBenchmarkingParams,
-    rng_or_seed: np.random.Generator | int,
     qubits: Optional[Sequence[ops.Qid] | Sequence[Sequence[ops.Qid]]] = None,
+    rng_or_seed: Optional[np.random.Generator | int] = None,
 ) -> tuple[Sequence[ResultDict], dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]:
     """Run the circuits in a shuffled order with readout error benchmarking.
 
@@ -459,8 +466,8 @@ def run_sweep_with_readout_benchmarking(
     input_circuits: list[circuits.Circuit],
     sweep_params: Sequence[study.Sweepable],
     parameters: ReadoutBenchmarkingParams,
-    rng_or_seed: np.random.Generator | int,
     qubits: Optional[Sequence[ops.Qid] | Sequence[Sequence[ops.Qid]]] = None,
+    rng_or_seed: Optional[np.random.Generator | int] = None,
 ) -> tuple[
     Sequence[Sequence[study.Result]], dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]
 ]:
@@ -470,11 +477,11 @@ def run_sweep_with_readout_benchmarking(
         input_circuits: The circuits to run.
         sweep_params: The sweep parameters for the input circuits.
         parameters: The readout benchmarking parameters.
-        rng_or_seed: A random number generator used to generate readout circuits.
-                     Or an integer seed.
         qubits: The qubits to benchmark readout errors. If None, all qubits in the
         input_circuits are used. Can be a list of qubits or a list of tuples
         of qubits.
+        rng_or_seed: A random number generator used to generate readout circuits.
+                     Or an integer seed.
 
     Returns:
         A tuple containing:
@@ -524,10 +531,9 @@ def run_sweep_with_readout_benchmarking(
 
     # Analyze results
     readout_calibration_results = {}
-    i = 0
-    for qubit_group, random_bitstrings in zip(qubits_to_measure, all_random_bitstrings):
-        group_measurements = readout_measurements[i]
-        i += 1
+    for qubit_group, random_bitstrings, group_measurements in zip(
+        qubits_to_measure, all_random_bitstrings, readout_measurements
+    ):
 
         calibration_result = _analyze_readout_results(
             group_measurements,
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
index fbb98ac895f..b4ef98c7273 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
@@ -94,7 +94,7 @@ def _circuits_with_readout_benchmarking_errors_shuffled(
 ):
     measurements, readout_calibration_results = (
         sc_readout.run_shuffled_circuits_with_readout_benchmarking(
-            sampler, input_circuits, parameters, rng_or_seed, qubits
+            sampler, input_circuits, parameters, qubits, rng_or_seed
         )
     )
 
@@ -113,7 +113,7 @@ def _circuits_with_readout_benchmarking_errors_sweep(
 ):
     sweep_measurements, readout_calibration_results = (
         sc_readout.run_sweep_with_readout_benchmarking(
-            sampler, input_circuits, sweep_params, parameters, rng_or_seed, qubits
+            sampler, input_circuits, sweep_params, parameters, qubits, rng_or_seed
         )
     )
 
@@ -345,12 +345,11 @@ def test_can_handle_zero_random_bitstring(mode: str):
 
 
 @pytest.mark.parametrize("mode", ["shuffled", "sweep"])
-def test_circuits_with_readout_benchmarking_no_qubits_arg(mode: str):
+def test_circuits_with_readout_benchmarking_no_qubits_arg_empty_rng(mode: str):
     """Test benchmarking when the `qubits` argument is not provided."""
     qubits = cirq.LineQubit.range(3)
     sampler = NoisySingleQubitReadoutSampler(p0=0.1, p1=0.2, seed=1234)
     circuit_repetitions = 1
-    rng = np.random.default_rng()
     readout_repetitions = 1000
     num_random_bitstrings = 100
 
@@ -364,7 +363,7 @@ def test_circuits_with_readout_benchmarking_no_qubits_arg(mode: str):
         input_circuits = _create_test_circuits(qubits, 3)
         measurements, readout_calibration_results = (
             sc_readout.run_shuffled_circuits_with_readout_benchmarking(
-                sampler, input_circuits, readout_benchmarking_params, rng, qubits=None
+                sampler, input_circuits, readout_benchmarking_params, None, None
             )
         )
         assert len(measurements) == len(input_circuits)
@@ -372,7 +371,7 @@ def test_circuits_with_readout_benchmarking_no_qubits_arg(mode: str):
         input_circuits, sweep_params = _create_test_circuits_with_sweep(qubits, 3)
         sweep_measurements, readout_calibration_results = (
             sc_readout.run_sweep_with_readout_benchmarking(
-                sampler, input_circuits, sweep_params, readout_benchmarking_params, rng, qubits=None
+                sampler, input_circuits, sweep_params, readout_benchmarking_params, None, None
             )
         )
         assert len(sweep_measurements) == len(input_circuits)

From 4c0a160e5e7ffa829df413c155d7cabdb7793b96 Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Tue, 29 Jul 2025 14:40:54 -0700
Subject: [PATCH 7/8] Add test coverage

---
 ...circuits_with_readout_benchmarking_test.py | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
index b4ef98c7273..c64c5dee64a 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
@@ -431,6 +431,44 @@ def test_deprecated_run_shuffled_with_readout_benchmarking():
     qlist_none, _ = list(results_none.items())[0]
     assert set(qlist_none) == set(qubits)
 
+    # Test circuit_repetitions must be > 0
+    with cirq.testing.assert_deprecated(deadline="v2.0", count=1):
+        with pytest.raises(ValueError, match="Must provide non-zero circuit_repetitions."):
+            sc_readout.run_shuffled_with_readout_benchmarking(
+                input_circuits,
+                sampler,
+                circuit_repetitions=0,
+                num_random_bitstrings=5,
+                readout_repetitions=100,
+                rng_or_seed=123,
+            )
+
+    # Test num_random_bitstrings must be >= 0
+    with cirq.testing.assert_deprecated(deadline="v2.0", count=1):
+        with pytest.raises(ValueError, match="Must provide zero or more num_random_bitstrings."):
+            sc_readout.run_shuffled_with_readout_benchmarking(
+                input_circuits,
+                sampler,
+                circuit_repetitions=10,
+                num_random_bitstrings=-1,
+                readout_repetitions=100,
+                rng_or_seed=123,
+            )
+
+    # Test readout_repetitions must be > 0
+    with cirq.testing.assert_deprecated(deadline="v2.0", count=1):
+        with pytest.raises(
+            ValueError, match="Must provide non-zero readout_repetitions for readout calibration."
+        ):
+            sc_readout.run_shuffled_with_readout_benchmarking(
+                input_circuits,
+                sampler,
+                circuit_repetitions=10,
+                num_random_bitstrings=1,
+                readout_repetitions=0,
+                rng_or_seed=123,
+            )
+
 
 def test_empty_input_circuits():
     """Test that the input circuits are empty."""

From f175cc4e484b6abd77bff96862d5e0325c1f9b5d Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Tue, 29 Jul 2025 16:55:46 -0700
Subject: [PATCH 8/8] Remove the rng checks

---
 ...shuffle_circuits_with_readout_benchmarking.py |  8 --------
 ...le_circuits_with_readout_benchmarking_test.py | 16 ----------------
 2 files changed, 24 deletions(-)

diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
index 66d6b0c4245..033e79acc8d 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking.py
@@ -79,14 +79,6 @@ def _validate_experiment_input(
     if isinstance(circuit_repetitions, list) and len(circuit_repetitions) != len(input_circuits):
         raise ValueError("Number of circuit_repetitions must match the number of input circuits.")
 
-    # Check rng is a numpy random generator
-    if (
-        rng_or_seed
-        and not isinstance(rng_or_seed, np.random.Generator)
-        and not isinstance(rng_or_seed, int)
-    ):
-        raise ValueError("Must provide a numpy random generator or a seed")
-
 
 def _validate_experiment_input_with_sweep(
     input_circuits: Sequence[circuits.Circuit],
diff --git a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
index c64c5dee64a..7c20f677746 100644
--- a/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
+++ b/cirq-core/cirq/contrib/shuffle_circuits/shuffle_circuits_with_readout_benchmarking_test.py
@@ -560,22 +560,6 @@ def test_zero_readout_repetitions():
         )
 
 
-def test_rng_type_mismatch():
-    """Test that the rng is not a numpy random generator or a seed."""
-    q = cirq.LineQubit(0)
-    circuit = cirq.Circuit(cirq.H(q), cirq.measure(q))
-    readout_benchmarking_params = sc_readout.ReadoutBenchmarkingParams(
-        circuit_repetitions=10, num_random_bitstrings=5, readout_repetitions=100
-    )
-    with pytest.raises(ValueError, match="Must provide a numpy random generator or a seed"):
-        sc_readout.run_shuffled_circuits_with_readout_benchmarking(
-            cirq.ZerosSampler(),
-            [circuit],
-            readout_benchmarking_params,
-            rng_or_seed="not a random generator or seed",
-        )
-
-
 def test_empty_sweep_params():
     """Test that the sweep params are empty."""
     q = cirq.LineQubit(5)