From b386c35a4138b389cc5487a20ea0e7a23b1d0cee Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Mon, 23 Mar 2026 20:18:55 -0700
Subject: [PATCH 01/11] (2/2) Combine Pauli measurements and postselection

---
 ...ing_measurement_with_readout_mitigation.py | 380 ++++++++++++++++--
 1 file changed, 341 insertions(+), 39 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 28f85694850..17c390cdca2 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
@@ -19,7 +19,7 @@
 import itertools
 import time
 from collections.abc import Mapping, Sequence
-from typing import cast, TYPE_CHECKING
+from typing import cast, TYPE_CHECKING, Union
 
 import attrs
 import numpy as np
@@ -377,50 +377,76 @@ def _validate_and_normalize_unformatted_input(
     return param_list
 
 
-def _extract_readout_qubits(pauli_strings: Sequence[ops.PauliString]) -> list[ops.Qid]:
-    """Extracts unique qubits from a list of QWC Pauli strings."""
-    return sorted({q for ps in pauli_strings for q in ps.qubits})
-
-
-def _pauli_strings_to_basis_change_ops(
-    pauli_strings: Sequence[ops.PauliString], qid_list: Sequence[ops.Qid]
+def _extract_readout_qubits(pauli_strings: Sequence[ops.PauliString], symmetries: list[ops.PauliString | ops.PauliSum] | None = None) -> list[ops.Qid]:
+    """Extracts unique qubits from both the target Pauli strings and the symmetries."""
+    all_qubits = set(q for ps in pauli_strings for q in ps.qubits)
+    
+    if symmetries:
+        for sym in symmetries:
+            if isinstance(sym, ops.PauliString):
+                all_qubits.update(sym.qubits)
+            elif isinstance(sym, ops.PauliSum):
+                for term in sym:
+                    all_qubits.update(term.qubits)
+                    
+    return sorted(list(all_qubits))
+    
+
+def _pauli_objs_to_basis_change_ops(
+    pauli_objs: Sequence[Union[ops.PauliString, ops.PauliSum]], qid_list: Sequence[ops.Qid]
 ):
+    flattened_terms = []
+    for obj in pauli_objs:
+        if isinstance(obj, ops.PauliString):
+            flattened_terms.append(obj)
+        elif isinstance(obj, ops.PauliSum):
+            flattened_terms.extend(list(obj))
+
     operations = []
+    
     for qubit in qid_list:
-        for pauli_str in pauli_strings:
+        for pauli_str in flattened_terms:
             pauli_op = pauli_str.get(qubit, default=ops.I)
             if pauli_op == ops.X:
-                operations.append(ops.ry(-np.pi / 2)(qubit))  # =cirq.H
+                operations.append(ops.ry(-np.pi / 2)(qubit))
                 break
-            if pauli_op == ops.Y:
+            elif pauli_op == ops.Y:
                 operations.append(ops.rx(np.pi / 2)(qubit))
                 break
-    return operations
+                
+    return operations      
 
 
-def _pauli_strings_to_basis_change_with_sweep(
-    pauli_strings: Sequence[ops.PauliString], qid_list: Sequence[ops.Qid]
+def _pauli_objs_to_basis_change_with_sweep(
+    pauli_objs: Sequence[Union[ops.PauliString| ops.PauliSum]], qid_list: Sequence[ops.Qid]
 ) -> dict[str, float]:
     """Decide single-qubit rotation sweep parameters for basis change.
 
     Args:
-        pauli_strings: A list of QWC Pauli strings.
+        pauli_objects: A list of QWC Pauli strings or Pauli sums.
         qid_list: A list of qubits to apply the basis change on.
     Returns:
         A dictionary mapping parameter names to their values for basis change.
     """
-    params_dict = {}
+    flattened_terms = []
+    for obj in pauli_objs:
+        if isinstance(obj, ops.PauliString):
+            flattened_terms.append(obj)
+        elif isinstance(obj, ops.PauliSum):
+            flattened_terms.extend(list(obj))
+
+    params_dict: dict[str, float] = {}
 
     for qid, qubit in enumerate(qid_list):
         params_dict[f"phi{qid}"] = 1.0
         params_dict[f"theta{qid}"] = 0.0
-        for pauli_str in pauli_strings:
+        for pauli_str in flattened_terms:
             pauli_op = pauli_str.get(qubit, default=ops.I)
             if pauli_op == ops.X:
                 params_dict[f"phi{qid}"] = 0.0
                 params_dict[f"theta{qid}"] = 1 / 2
                 break
-            if pauli_op == ops.Y:
+            elif pauli_op == ops.Y:
                 params_dict[f"phi{qid}"] = 1.0
                 params_dict[f"theta{qid}"] = 1 / 2
                 break
@@ -438,6 +464,7 @@ def _generate_basis_change_circuits(
     for params in normalized_circuits_to_pauli:
         input_circuit = params.circuit
         pauli_string_groups = params.pauli_strings
+        symmetries = tuple(sym for sym, _ in params.postselection_symmetries)
 
         global_qubits = list(qubits_to_measure) if qubits_to_measure is not None else None
         basis_change_circuits = []
@@ -448,11 +475,11 @@ def _generate_basis_change_circuits(
                 current_qid_list = global_qubits
             else:
                 # Extract qubits from Pauli strings
-                current_qid_list = _extract_readout_qubits(pauli_strings)
+                current_qid_list = _extract_readout_qubits(pauli_strings, symmetries)
 
             basis_change_circuit = circuits.Circuit(
                 input_circuit_unfrozen,
-                _pauli_strings_to_basis_change_ops(pauli_strings, current_qid_list),
+                _pauli_objs_to_basis_change_ops(pauli_strings + symmetries, current_qid_list),
                 ops.measure(*current_qid_list, key="result"),
                 strategy=insert_strategy,
             )
@@ -473,6 +500,8 @@ def _generate_basis_change_circuits_with_sweep(
     for params in normalized_circuits_to_pauli:
         input_circuit = params.circuit
         pauli_string_groups = params.pauli_strings
+        symmetries = tuple(sym for sym, _ in params.postselection_symmetries)
+
 
         # If qubits_to_measure is provided, use it
         if qubits_to_measure:
@@ -492,7 +521,7 @@ def _generate_basis_change_circuits_with_sweep(
             sweep_param = []
             for pauli_strings in pauli_string_groups:
                 sweep_param.append(
-                    _pauli_strings_to_basis_change_with_sweep(pauli_strings, qubits_to_measure)
+                    _pauli_objs_to_basis_change_with_sweep(pauli_strings, qubits_to_measure)
                 )
             sweep_params.append(sweep_param)
             parameterized_circuits.append(parameterized_circuit)
@@ -500,7 +529,7 @@ def _generate_basis_change_circuits_with_sweep(
         else:
             for pauli_strings in pauli_string_groups:
                 # Extract qubits from Pauli strings
-                qid_list = _extract_readout_qubits(pauli_strings)
+                qid_list = _extract_readout_qubits(pauli_strings, symmetries)
                 phi_symbols = sympy.symbols(f"phi:{len(qid_list)}")
                 theta_symbols = sympy.symbols(f"theta:{len(qid_list)}")
                 # Create phased gates and measurement operator
@@ -512,8 +541,8 @@ def _generate_basis_change_circuits_with_sweep(
                 parameterized_circuit = circuits.Circuit(
                     input_circuit.unfreeze(), phased_gates, measurement_op, strategy=insert_strategy
                 )
-                sweep_param_dict = _pauli_strings_to_basis_change_with_sweep(
-                    pauli_strings, qid_list
+                sweep_param_dict = _pauli_objs_to_basis_change_with_sweep(
+                    pauli_strings + symmetries, qid_list
                 )
                 parameterized_circuits.append(parameterized_circuit)
                 sweep_params.append(sweep_param_dict)
@@ -565,6 +594,101 @@ def _build_many_one_qubits_empty_confusion_matrix(qubits_length: int) -> list[np
     return [_build_one_qubit_confusion_matrix(0, 0) for _ in range(qubits_length)]
 
 
+def _split_input_circuits(
+    circuits_to_pauli_params: list[CircuitToPauliStringsParameters],
+) -> tuple[list[CircuitToPauliStringsParameters], list[CircuitToPauliStringsParameters]]:
+    """Splits the input circuits into two lists based on the way they are measured."""
+    # Circuits could be measured based on symmetries
+    symmetry_circuits: list[CircuitToPauliStringsParameters] = []
+    # Circuits could be measured based on confusion matrices
+    confusion_circuits: list[CircuitToPauliStringsParameters] = []
+
+    for circuit_to_pauli_params in circuits_to_pauli_params:
+        if not circuit_to_pauli_params.postselection_symmetries:
+            # If no postselection symmetries are provided, treat the circuit as a confusion circuit
+            confusion_circuits.append(circuit_to_pauli_params)
+            continue
+        else:
+            symmetry_circuits.append(circuit_to_pauli_params)
+    return symmetry_circuits, confusion_circuits
+
+
+def _process_symmetry_measurement_results(
+    qubits: Sequence[ops.Qid],
+    pauli_string_groups: list[list[ops.PauliString]],
+    measurement_results: np.ndarray,
+    circuit_to_pauli: CircuitToPauliStringsParameters,
+    pauli_repetitions: int,
+) -> list[PauliStringMeasurementResult]:
+    """Filters measurement results using symmetries and calculates expectations."""
+    single_circuit_pauli_measurement_results: list[PauliStringMeasurementResult] = []
+
+    # filter out bitstrings based on postselection symmetries
+    measurement_result_eigenvalues = 1 - 2 * measurement_results
+    rows_to_keep_mask = np.ones(len(measurement_result_eigenvalues), dtype=bool)
+
+    for sym, expected_value in circuit_to_pauli.postselection_symmetries:
+        # Determine which rows to keep based on the symmetry
+        if isinstance(sym, ops.PauliString):
+            sym_qubit_indices = [qubits.index(q) for q in sym.keys()]
+            actual_eigenvalues = np.prod(
+                measurement_result_eigenvalues[:, sym_qubit_indices], axis=1
+            )
+            rows_to_keep_mask &= actual_eigenvalues == expected_value
+
+        elif isinstance(sym, ops.PauliSum):
+            sum_eigenvalues = np.zeros(len(measurement_result_eigenvalues), dtype=float)
+            for term in sym:
+                term_qubit_indices = [qubits.index(q) for q in term.keys()]
+                term_eigenvalues = np.prod(
+                    measurement_result_eigenvalues[:, term_qubit_indices], axis=1
+                )
+                sum_eigenvalues += term.coefficient.real * term_eigenvalues
+            rows_to_keep_mask &= np.isclose(sum_eigenvalues, expected_value)
+
+    post_selection_circuits_results = measurement_results[rows_to_keep_mask]
+
+    for pauli_str in pauli_string_groups:
+        qubits_sorted = sorted(pauli_str.qubits)
+        qubit_indices = [qubits.index(q) for q in qubits_sorted]
+        relevant_bits_unmit = measurement_results[:, qubit_indices]
+
+        if len(post_selection_circuits_results) == 0:
+            raw_mitigated_values = np.nan
+            raw_d_m = np.nan
+        else:
+            relevant_bits_mit = post_selection_circuits_results[:, qubit_indices]
+            parity = np.sum(relevant_bits_mit, axis=1) % 2
+            raw_mitigated_values = 1 - 2 * np.mean(parity)
+            raw_d_m = 2 * np.sqrt(np.mean(parity) * (1 - np.mean(parity)) / len(relevant_bits_mit))
+
+        mitigated_value_with_coefficient = raw_mitigated_values * pauli_str.coefficient.real
+        d_mit_with_coefficient = raw_d_m * abs(pauli_str.coefficient.real)
+
+        # Calculate the unmitigated expectation.
+        parity_unmit = np.sum(relevant_bits_unmit, axis=1) % 2
+        raw_unmitigated_values = 1 - 2 * np.mean(parity_unmit)
+        raw_d_unmit = 2 * np.sqrt(
+            np.mean(parity_unmit) * (1 - np.mean(parity_unmit)) / pauli_repetitions
+        )
+        unmitigated_value_with_coefficient = raw_unmitigated_values * pauli_str.coefficient.real
+        d_unmit_with_coefficient = raw_d_unmit * abs(pauli_str.coefficient.real)
+
+        single_circuit_pauli_measurement_results.append(
+            PauliStringMeasurementResult(
+                pauli_string=pauli_str,
+                mitigated_expectation=mitigated_value_with_coefficient,
+                mitigated_stddev=d_mit_with_coefficient,
+                unmitigated_expectation=unmitigated_value_with_coefficient,
+                unmitigated_stddev=d_unmit_with_coefficient,
+                calibration_result=PostFilteringSymmetryCalibrationResult(
+                    raw_bitstrings=measurement_results,
+                    filtered_bitstrings=post_selection_circuits_results,
+                ),
+            )
+        )
+    return single_circuit_pauli_measurement_results
+
 def _process_pauli_measurement_results(
     pauli_string_groups: Sequence[Sequence[ops.PauliString]],
     circuit_results: Sequence[cirq.ResultDict] | Sequence[cirq.Result],
@@ -678,13 +802,123 @@ def _process_pauli_measurement_results(
 
     return pauli_measurement_results
 
+def _measure_pauli_strings_with_symmetries(
+    sampler: work.Sampler,
+    circuits_to_pauli: list[CircuitToPauliStringsParameters],
+    pauli_repetitions: int,
+    use_sweep: bool,
+    insert_strategy: circuits.InsertStrategy,
+) -> list[CircuitToPauliStringsMeasurementResult]:
+    """
+    Measures expectation values of Pauli strings on given circuits with postselection symmetries.
+    This function takes a list of CircuitToPauliStringsParameters. Each parameter contains a circuit, its associated list of QWC Pauli string groups and postselection symmetries.
+    For each circuit_to_pauli, it:
+    1. Runs the circuits to get the measurement results.
+    2. Filters the measurement results based on postselection symmetries.
+    3. Calculates and returns the expectation values for each Pauli string.
 
-def measure_pauli_strings(
-    circuits_to_pauli: (
-        Mapping[circuits.FrozenCircuit, Sequence[ops.PauliString]]
-        | Mapping[circuits.FrozenCircuit, Sequence[Sequence[ops.PauliString]]]
-        | list[CircuitToPauliStringsParameters]
-    ),
+    Args:
+        sampler: The sampler to use.
+        circuits_to_pauli_params: A list of CircuitToPauliStringsParameters objects, where
+            each object contains:
+            - The circuit to measure.
+            - A list of Pauli strings or a list of lists of QWC Pauli strings.
+            - A dictionary mapping Pauli strings or Pauli sums to expected eigen value for
+              postselection symmetries.
+        pauli_repetitions: The number of repetitions for each circuit when measuring
+            Pauli strings.
+        use_sweep: Whether to use parameter sweeps for basis change circuits.
+        insert_strategy: The strategy to use when inserting basis change and measurement
+    """
+    # Skip if no circuits to measure
+    if not circuits_to_pauli:
+        return []
+    
+    final_measurement_results: list[CircuitToPauliStringsMeasurementResult] = []
+
+    # Generate measurement circuits
+    if use_sweep:
+        pauli_measurement_circuits, sweep_params = _generate_basis_change_circuits_with_sweep(
+            circuits_to_pauli, insert_strategy
+        )
+
+        # Run the sweeps
+        all_circuits_measurement_results: list[list[study.Result]] = []
+        for parameterized_circuit, sweep in zip(pauli_measurement_circuits, sweep_params):
+            results_for_one_circuit = sampler.run_sweep(
+                program=parameterized_circuit, params=sweep, repetitions=pauli_repetitions
+            )
+            all_circuits_measurement_results.append(results_for_one_circuit)
+
+        # Process the results
+        for circuit_to_pauli_params, circuit_results in zip(
+            circuits_to_pauli, all_circuits_measurement_results
+        ):
+            qubits_in_circuit = tuple(sorted(circuit_to_pauli_params.circuit.all_qubits()))
+            single_circuit_pauli_measurement_results: list[PauliStringMeasurementResult] = []
+
+            for i, circuit_result in enumerate(circuit_results):
+                single_circuit_pauli_measurement_results.extend(
+                    _process_symmetry_measurement_results(
+                        qubits_in_circuit,
+                        circuit_to_pauli_params.pauli_strings[i],
+                        circuit_result.measurements["result"],
+                        circuit_to_pauli_params,
+                        pauli_repetitions,
+                    )
+                )
+
+            final_measurement_results.append(
+                CircuitToPauliStringsMeasurementResult(
+                    circuit=circuit_to_pauli_params.circuit,
+                    results=single_circuit_pauli_measurement_results,
+                )
+            )
+    else:
+        # Process in batch mode
+        pauli_measurement_circuits = _generate_basis_change_circuits(
+            circuits_to_pauli, insert_strategy
+        )
+
+        circuits_results = sampler.run_batch(
+            pauli_measurement_circuits, repetitions=pauli_repetitions
+        )
+        circuits_measurement_results = [cir[0] for cir in circuits_results]
+
+        circuit_result_index = 0
+        for circuit_to_pauli_params in circuits_to_pauli:
+            circuit_results = circuits_measurement_results[
+                circuit_result_index : circuit_result_index
+                + len(circuit_to_pauli_params.pauli_strings)
+            ]
+            qubits_in_circuit = tuple(sorted(circuit_to_pauli_params.circuit.all_qubits()))
+            single_circuit_pauli_measurement_results: list[PauliStringMeasurementResult] = []
+
+            for i, circuit_result in enumerate(circuit_results):
+                single_circuit_pauli_measurement_results.extend(
+                    _process_symmetry_measurement_results(
+                        qubits_in_circuit,
+                        circuit_to_pauli_params.pauli_strings[i],
+                        circuit_result.measurements["result"],
+                        circuit_to_pauli_params,
+                        pauli_repetitions,
+                    )
+                )
+
+            circuit_result_index += len(circuit_to_pauli_params.pauli_strings)
+            final_measurement_results.append(
+                CircuitToPauliStringsMeasurementResult(
+                    circuit=circuit_to_pauli_params.circuit,
+                    results=single_circuit_pauli_measurement_results,
+                )
+            )
+    return final_measurement_results
+
+    
+
+
+def _measure_pauli_strings_with_confusion_matrices(
+    normalized_circuits_to_pauli: list[CircuitToPauliStringsParameters],
     sampler: work.Sampler,
     pauli_repetitions: int,
     readout_repetitions: int,
@@ -695,7 +929,7 @@ def measure_pauli_strings(
     measure_on_full_support: bool = False,
 ) -> list[CircuitToPauliStringsMeasurementResult]:
     """Measures expectation values of Pauli strings on given circuits with/without
-    readout error mitigation.
+    readout error mitigation based on confusion matrices.
 
     For each circuit and its associated list of QWC pauli string group, it:
     1.  Constructs circuits to measure the Pauli string expectation value by
@@ -741,14 +975,8 @@ def measure_pauli_strings(
             - A list of PauliStringMeasurementResult objects.
             - The calibration result for single-qubit readout errors.
     """
-
-    normalized_circuits_to_pauli = _validate_input(
-        circuits_to_pauli,
-        pauli_repetitions,
-        readout_repetitions,
-        num_random_bitstrings,
-        rng_or_seed,
-    )
+    if not normalized_circuits_to_pauli:
+        return []
 
     # Extract unique qubit tuples from input pauli strings
     if measure_on_full_support:
@@ -866,3 +1094,77 @@ def measure_pauli_strings(
         )
 
     return results
+
+
+def measure_pauli_strings(
+    circuits_to_pauli: (
+        Mapping[circuits.FrozenCircuit, Sequence[ops.PauliString]]
+        | Mapping[circuits.FrozenCircuit, Sequence[Sequence[ops.PauliString]]]
+        | list[CircuitToPauliStringsParameters]
+    ),
+    sampler: work.Sampler,
+    pauli_repetitions: int,
+    readout_repetitions: int,
+    num_random_bitstrings: int,
+    rng_or_seed: np.random.Generator | int,
+    use_sweep: bool = False,
+    insert_strategy: circuits.InsertStrategy = circuits.InsertStrategy.INLINE,
+    measure_on_full_support: bool = False,
+) -> list[CircuitToPauliStringsMeasurementResult]:
+    """Measures expectation values of Pauli strings on given circuits with/without
+    readout error mitigation.
+
+    Args:
+        circuits_to_pauli: A list of CircuitToPauliStringsParameters objects, where each object contains:
+            - The circuit to measure.
+            - A list of QWC groups (list[list[ops.PauliString]]) or a list of PauliStrings (list[ops.PauliString]).
+            - A dictionary mapping Pauli strings or Pauli sums to expected eigen value for postselection symmetries.
+        sampler: The sampler to use.
+        pauli_repetitions: The number of repetitions for each circuit when measuring
+            Pauli strings.
+        readout_repetitions: The number of repetitions for readout calibration
+            in the shuffled benchmarking.
+        num_random_bitstrings: The number of random bitstrings to use in readout
+            benchmarking.
+        rng_or_seed: A random number generator or seed for the readout benchmarking.
+        use_sweep: If True, uses parameterized circuits and sweeps parameters
+            for both Pauli measurements and readout benchmarking. Defaults to False.
+        insert_strategy: The strategy for inserting measurement operations into the circuit.
+            Defaults to circuits.InsertStrategy.INLINE.
+
+    Returns:
+        A list of CircuitToPauliStringsMeasurementResult objects, where each object contains:
+            - The circuit that was measured.
+            - A list of PauliStringMeasurementResult objects.
+            - The calibration result for single-qubit readout errors.
+    """
+     
+    normalized_circuits_to_pauli = _validate_input(
+        circuits_to_pauli,
+        pauli_repetitions,
+        readout_repetitions,
+        num_random_bitstrings,
+        rng_or_seed,
+    )
+
+
+    # Split the input circuits into two lists based on the way they are measured.
+    symmetry_circuits, confusion_circuits = _split_input_circuits(normalized_circuits_to_pauli)
+
+    return _measure_pauli_strings_with_symmetries(
+        sampler=sampler,
+        circuits_to_pauli=symmetry_circuits,
+        pauli_repetitions=pauli_repetitions,
+        use_sweep=use_sweep,
+        insert_strategy=insert_strategy,
+    ) + _measure_pauli_strings_with_confusion_matrices(
+        sampler=sampler,
+       normalized_circuits_to_pauli=confusion_circuits,
+        pauli_repetitions=pauli_repetitions,
+        readout_repetitions=readout_repetitions,
+        num_random_bitstrings=num_random_bitstrings,
+        rng_or_seed=rng_or_seed,
+        use_sweep=use_sweep,
+        insert_strategy=insert_strategy,
+        measure_on_full_support = measure_on_full_support
+    )
\ No newline at end of file

From 1fecad7c2f7a95e97f83efb99c280a512934f2d8 Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Wed, 25 Mar 2026 14:22:39 -0700
Subject: [PATCH 02/11] Add tests

---
 .../cirq/contrib/paulistring/__init__.py      |   1 +
 ...ing_measurement_with_readout_mitigation.py |  50 ++--
 ...easurement_with_readout_mitigation_test.py | 214 ++++++++++--------
 3 files changed, 137 insertions(+), 128 deletions(-)

diff --git a/cirq-core/cirq/contrib/paulistring/__init__.py b/cirq-core/cirq/contrib/paulistring/__init__.py
index 782f5cb8ed2..cd8481eb22b 100644
--- a/cirq-core/cirq/contrib/paulistring/__init__.py
+++ b/cirq-core/cirq/contrib/paulistring/__init__.py
@@ -46,4 +46,5 @@
 from cirq.contrib.paulistring.pauli_string_measurement_with_readout_mitigation import (
     measure_pauli_strings as measure_pauli_strings,
     CircuitToPauliStringsParameters as CircuitToPauliStringsParameters,
+    PostFilteringSymmetryCalibrationResult as PostFilteringSymmetryCalibrationResult,
 )
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 17c390cdca2..a5d56c89d42 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
@@ -243,11 +243,6 @@ def _validate_circuit_to_pauli_strings_parameters(
                 _validate_single_pauli_string(pauli_str)
 
         # 3. Validate postselection symmetries
-        # Postselection symmetries are parsed and validated below, however
-        # the functionality to actually apply postselection filtering is not
-        # implemented in this PR. Raise a clear error if the user attempts to
-        # provide postselection symmetries so users do not assume they are
-        # applied.
         for sym, _ in params.postselection_symmetries:
             if isinstance(sym, ops.PauliSum):
                 terms = tuple(sym)
@@ -278,14 +273,6 @@ def _validate_circuit_to_pauli_strings_parameters(
             raise ValueError(
                 f"Postselection symmetries of {params.circuit} are not commuting with all Pauli"
             )
-        # Postselection symmetries are parsed and validated above, however
-        # the functionality to actually apply postselection filtering is not
-        # implemented in this PR.
-        if params.postselection_symmetries:
-            raise NotImplementedError(
-                "Postselection symmetries are not implemented. "
-                "Please use readout mitigation via confusion matrices instead."
-            )
 
 
 def _validate_input(
@@ -377,10 +364,13 @@ def _validate_and_normalize_unformatted_input(
     return param_list
 
 
-def _extract_readout_qubits(pauli_strings: Sequence[ops.PauliString], symmetries: list[ops.PauliString | ops.PauliSum] | None = None) -> list[ops.Qid]:
+def _extract_readout_qubits(
+    pauli_strings: Sequence[ops.PauliString],
+    symmetries: Sequence[ops.PauliString | ops.PauliSum] | None = None,
+) -> list[ops.Qid]:
     """Extracts unique qubits from both the target Pauli strings and the symmetries."""
     all_qubits = set(q for ps in pauli_strings for q in ps.qubits)
-    
+
     if symmetries:
         for sym in symmetries:
             if isinstance(sym, ops.PauliString):
@@ -388,9 +378,9 @@ def _extract_readout_qubits(pauli_strings: Sequence[ops.PauliString], symmetries
             elif isinstance(sym, ops.PauliSum):
                 for term in sym:
                     all_qubits.update(term.qubits)
-                    
+
     return sorted(list(all_qubits))
-    
+
 
 def _pauli_objs_to_basis_change_ops(
     pauli_objs: Sequence[Union[ops.PauliString, ops.PauliSum]], qid_list: Sequence[ops.Qid]
@@ -403,7 +393,7 @@ def _pauli_objs_to_basis_change_ops(
             flattened_terms.extend(list(obj))
 
     operations = []
-    
+
     for qubit in qid_list:
         for pauli_str in flattened_terms:
             pauli_op = pauli_str.get(qubit, default=ops.I)
@@ -413,12 +403,12 @@ def _pauli_objs_to_basis_change_ops(
             elif pauli_op == ops.Y:
                 operations.append(ops.rx(np.pi / 2)(qubit))
                 break
-                
-    return operations      
+
+    return operations
 
 
 def _pauli_objs_to_basis_change_with_sweep(
-    pauli_objs: Sequence[Union[ops.PauliString| ops.PauliSum]], qid_list: Sequence[ops.Qid]
+    pauli_objs: Sequence[Union[ops.PauliString | ops.PauliSum]], qid_list: Sequence[ops.Qid]
 ) -> dict[str, float]:
     """Decide single-qubit rotation sweep parameters for basis change.
 
@@ -502,7 +492,6 @@ def _generate_basis_change_circuits_with_sweep(
         pauli_string_groups = params.pauli_strings
         symmetries = tuple(sym for sym, _ in params.postselection_symmetries)
 
-
         # If qubits_to_measure is provided, use it
         if qubits_to_measure:
             phi_symbols = sympy.symbols(f"phi:{len(qubits_to_measure)}")
@@ -615,7 +604,7 @@ def _split_input_circuits(
 
 def _process_symmetry_measurement_results(
     qubits: Sequence[ops.Qid],
-    pauli_string_groups: list[list[ops.PauliString]],
+    pauli_string_groups: Sequence[ops.PauliString],
     measurement_results: np.ndarray,
     circuit_to_pauli: CircuitToPauliStringsParameters,
     pauli_repetitions: int,
@@ -689,6 +678,7 @@ def _process_symmetry_measurement_results(
         )
     return single_circuit_pauli_measurement_results
 
+
 def _process_pauli_measurement_results(
     pauli_string_groups: Sequence[Sequence[ops.PauliString]],
     circuit_results: Sequence[cirq.ResultDict] | Sequence[cirq.Result],
@@ -802,6 +792,7 @@ def _process_pauli_measurement_results(
 
     return pauli_measurement_results
 
+
 def _measure_pauli_strings_with_symmetries(
     sampler: work.Sampler,
     circuits_to_pauli: list[CircuitToPauliStringsParameters],
@@ -833,7 +824,7 @@ def _measure_pauli_strings_with_symmetries(
     # Skip if no circuits to measure
     if not circuits_to_pauli:
         return []
-    
+
     final_measurement_results: list[CircuitToPauliStringsMeasurementResult] = []
 
     # Generate measurement circuits
@@ -914,8 +905,6 @@ def _measure_pauli_strings_with_symmetries(
             )
     return final_measurement_results
 
-    
-
 
 def _measure_pauli_strings_with_confusion_matrices(
     normalized_circuits_to_pauli: list[CircuitToPauliStringsParameters],
@@ -1138,7 +1127,7 @@ def measure_pauli_strings(
             - A list of PauliStringMeasurementResult objects.
             - The calibration result for single-qubit readout errors.
     """
-     
+
     normalized_circuits_to_pauli = _validate_input(
         circuits_to_pauli,
         pauli_repetitions,
@@ -1147,7 +1136,6 @@ def measure_pauli_strings(
         rng_or_seed,
     )
 
-
     # Split the input circuits into two lists based on the way they are measured.
     symmetry_circuits, confusion_circuits = _split_input_circuits(normalized_circuits_to_pauli)
 
@@ -1159,12 +1147,12 @@ def measure_pauli_strings(
         insert_strategy=insert_strategy,
     ) + _measure_pauli_strings_with_confusion_matrices(
         sampler=sampler,
-       normalized_circuits_to_pauli=confusion_circuits,
+        normalized_circuits_to_pauli=confusion_circuits,
         pauli_repetitions=pauli_repetitions,
         readout_repetitions=readout_repetitions,
         num_random_bitstrings=num_random_bitstrings,
         rng_or_seed=rng_or_seed,
         use_sweep=use_sweep,
         insert_strategy=insert_strategy,
-        measure_on_full_support = measure_on_full_support
-    )
\ No newline at end of file
+        measure_on_full_support=measure_on_full_support,
+    )
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 61d4bdd632f..846e662c842 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
@@ -22,7 +22,11 @@
 import pytest
 
 import cirq
-from cirq.contrib.paulistring import CircuitToPauliStringsParameters, measure_pauli_strings
+from cirq.contrib.paulistring import (
+    CircuitToPauliStringsParameters,
+    measure_pauli_strings,
+    PostFilteringSymmetryCalibrationResult,
+)
 from cirq.experiments import SingleQubitReadoutCalibrationResult
 from cirq.experiments.single_qubit_readout_calibration_test import NoisySingleQubitReadoutSampler
 
@@ -313,77 +317,67 @@ def test_group_pauli_string_measurement_errors_with_noise(use_sweep: bool) -> No
 
 
 @pytest.mark.parametrize("use_sweep", [True, False])
-def test_many_circuits_input_measurement_with_noise(use_sweep: bool) -> None:
-    """Test that the mitigated expectation is close to the ideal expectation
-    based on the Pauli string for multiple circuits"""
+def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
+    """Test mixed input: some circuits using confusion matrices, some using symmetries.
+    
+    This test specifically includes a QWC group with multiple Pauli strings to 
+    ensure the processing logic handles nested groups correctly.
+    """
     qubits_1 = cirq.LineQubit.range(3)
-    qubits_2 = [
-        cirq.GridQubit(0, 1),
-        cirq.GridQubit(1, 1),
-        cirq.GridQubit(1, 0),
-        cirq.GridQubit(1, 2),
-        cirq.GridQubit(2, 1),
-    ]
-    qubits_3 = cirq.LineQubit.range(8)
+    qubits_2 = cirq.LineQubit.range(5)
 
+    # Circuit 1 is a Standard GHZ without symmetry.
     circuit_1 = cirq.FrozenCircuit(_create_ghz(3, qubits_1))
-    circuit_2 = cirq.FrozenCircuit(_create_ghz(5, qubits_2))
-    circuit_3 = cirq.FrozenCircuit(_create_ghz(8, qubits_3))
+    params_1 = CircuitToPauliStringsParameters(
+        circuit=circuit_1, 
+        pauli_strings=((cirq.PauliString(cirq.Z(qubits_1[0])),),)
+    )
 
-    circuits_to_pauli: dict[
-        cirq.FrozenCircuit, Sequence[cirq.PauliString] | Sequence[Sequence[cirq.PauliString]]
-    ] = {}
-    # This is to test mixed types could be handled.
-    circuits_to_pauli[circuit_1] = [
-        _generate_qwc_paulis(
-            _generate_random_pauli_string(qubits_1, enable_coeff=True, allow_pauli_i=False), 6
-        )
-        for _ in range(3)
-    ]
-    circuits_to_pauli[circuit_2] = [_generate_random_pauli_string(qubits_2) for _ in range(3)]
-    circuits_to_pauli[circuit_3] = [_generate_random_pauli_string(qubits_3) for _ in range(3)]
+    # Circuit 2 is a GHZ with a known symmetry (Z0*Z1 = 1)
+    circuit_2 = cirq.FrozenCircuit(_create_ghz(5, qubits_2))
+    symmetry: cirq.PauliString = cirq.PauliString(cirq.Z(qubits_2[0]) * cirq.Z(qubits_2[1]))
 
-    sampler = NoisySingleQubitReadoutSampler(p0=0.03, p1=0.05, seed=1234)
-    simulator = cirq.Simulator()
+    pauli_group = (
+        cirq.PauliString(cirq.Z(qubits_2[0])),
+        2.0 * cirq.PauliString(cirq.Z(qubits_2[1]))
+    )
 
-    circuits_with_pauli_expectations = measure_pauli_strings(
-        circuits_to_pauli,  # type: ignore
-        sampler,
-        300,
-        300,
-        300,
-        np.random.default_rng(),
-        use_sweep,
+    params_sym = CircuitToPauliStringsParameters(
+        circuit=circuit_2,
+        pauli_strings=(pauli_group,), 
+        postselection_symmetries=((symmetry, 1),),
     )
 
-    for circuit_with_pauli_expectations in circuits_with_pauli_expectations:
-        assert isinstance(circuit_with_pauli_expectations.circuit, cirq.FrozenCircuit)
+    sampler = NoisySingleQubitReadoutSampler(p0=0.01, p1=0.02, seed=1234)
 
-        expected_val_simulation = simulator.simulate(
-            circuit_with_pauli_expectations.circuit.unfreeze()
-        )
-        final_state_vector = expected_val_simulation.final_state_vector
+    results = measure_pauli_strings(
+        [params_1, params_sym],
+        sampler,
+        pauli_repetitions=500,
+        readout_repetitions=500,
+        num_random_bitstrings=100,
+        rng_or_seed=1234,
+        use_sweep=use_sweep,
+    )
 
-        for pauli_string_measurement_results in circuit_with_pauli_expectations.results:
-            assert np.isclose(
-                pauli_string_measurement_results.mitigated_expectation,
-                _ideal_expectation_based_on_pauli_string(
-                    pauli_string_measurement_results.pauli_string, final_state_vector
-                ),
-                atol=10 * pauli_string_measurement_results.mitigated_stddev,
-            )
-            assert isinstance(
-                pauli_string_measurement_results.calibration_result,
-                SingleQubitReadoutCalibrationResult,
-            )
-            for (
-                error
-            ) in pauli_string_measurement_results.calibration_result.zero_state_errors.values():
-                assert 0.025 < error < 0.035
-            for (
-                error
-            ) in pauli_string_measurement_results.calibration_result.one_state_errors.values():
-                assert 0.045 < error < 0.055
+    assert len(results) == 2
+
+    for circuit_res in results:
+        # For Circuit 2, we expect exactly 2 results because the group had 2 Pauli strings
+        if circuit_res.circuit == circuit_2:
+            assert len(circuit_res.results) == 2
+        
+        for res in circuit_res.results:
+            if circuit_res.circuit == circuit_2:
+                assert isinstance(res.calibration_result, PostFilteringSymmetryCalibrationResult)
+                # Verify coefficient was applied (2.0 for the second Pauli string)
+                if res.pauli_string.coefficient == 2.0:
+                    # Expectation should be roughly 2.0 * 1.0 or 2.0 * -1.0 depending on state
+                    # for GHZ Z expectation is 0, so 2.0 * 0 = 0.
+                    assert -2.2 <= res.mitigated_expectation <= 2.2
+            else:
+                assert isinstance(res.calibration_result, SingleQubitReadoutCalibrationResult)
+                assert -1.1 <= res.mitigated_expectation <= 1.1
 
 
 @pytest.mark.parametrize("use_sweep", [True, False])
@@ -807,9 +801,9 @@ def test_postselection_symmetry_validation_and_logic() -> None:
     params_valid_sum = CircuitToPauliStringsParameters(
         circuit=circuit, pauli_strings=target_paulis, postselection_symmetries=good_symmetries
     )
-    # Postselection is not implemented; providing symmetries should raise.
-    with pytest.raises(NotImplementedError, match="Postselection symmetries are not implemented"):
-        measure_pauli_strings([params_valid_sum], sampler, 10, 10, 0, rng)
+    results = measure_pauli_strings([params_valid_sum], sampler, 10, 1, 0, rng)
+    assert len(results) == 1
+    assert len(results[0].results) == 1
 
     # Test PauliSum with Non-QWC Terms
     # X0 and Z0 do not commute. This is an invalid PauliSum *structure* for this context.
@@ -858,44 +852,70 @@ def test_sampler_receives_correct_circuits(use_sweep: bool) -> None:
 
     from cirq.study.result import ResultDict
 
-    # Create a circuit with 5 qubits
     qubits = cirq.LineQubit.range(5)
-
     circuit = cirq.FrozenCircuit(_create_ghz(5, qubits))
+    pauli_qubits = qubits[1:4]  # Q1, Q2, Q3
+    pauli_str = cirq.PauliString(dict.fromkeys(pauli_qubits, cirq.X))
 
-    # Create Pauli strings that only use qubits 1, 2, 3 (indices 1,2,3)
-    pauli_qubits = qubits[1:4]  # qubits 1,2,3
-    pauli_string: cirq.PauliString = cirq.PauliString(dict.fromkeys(pauli_qubits, cirq.X))
-
-    circuits_to_pauli: dict[cirq.FrozenCircuit, list[cirq.PauliString]] = {}
-    circuits_to_pauli[circuit] = [pauli_string]
-
-    # Mock the sampler
+    # Test standard Pauli String without Symmetries
+    circuits_to_pauli = {circuit: [pauli_str]}
     mock_sampler = MagicMock()
 
-    # Configure the mock sampler to return a valid ResultDict
-    mock_sampler.run.return_value = ResultDict(
-        params=cirq.ParamResolver({}), measurements={"result": np.array([[0, 1, 0]])}
-    )
-    # Configure the mock sampler to return valid results for run_batch
-    mock_sampler.run_batch.return_value = [
-        [ResultDict(params=cirq.ParamResolver({}), measurements={"result": np.array([[0, 1, 0]])})]
-    ]
+    mock_res = ResultDict(params=cirq.ParamResolver({}), measurements={"result": np.zeros((1, 3))})
+    mock_sampler.run_batch.return_value = [[mock_res]]
+    mock_sampler.run_sweep.return_value = [mock_res]
 
-    # Call measure_pauli_strings with the mock sampler
-    measure_pauli_strings(circuits_to_pauli, mock_sampler, 100, 100, 0, 1234, use_sweep=use_sweep)
+    measure_pauli_strings(circuits_to_pauli, mock_sampler, 10, 10, 0, 1234, use_sweep=use_sweep)
 
-    # Verify the sampler was called with the correct circuits
-    batches = [call.args[0] for call in mock_sampler.run_batch.call_args_list]
-    called_circuits = [circuit for batch in batches for circuit in batch]
+    # Determine which mock method to check
+    if use_sweep:
+        called_circuits = [call.kwargs['program'] for call in mock_sampler.run_sweep.call_args_list]
+    else:
+        called_circuits = [
+            c for batch in mock_sampler.run_batch.call_args_list for c in batch.args[0]
+        ]
 
     for called_circuit in called_circuits:
-        measured_qubits = set()
-        for op in called_circuit.all_operations():
-            if isinstance(op.gate, cirq.MeasurementGate):
-                measured_qubits.update(op.qubits)
-
-        # Ensure only the qubits in the Pauli string are measured
-        assert measured_qubits == set(
-            pauli_qubits
-        ), f"Expected measured qubits: {set(pauli_qubits)}, but found: {measured_qubits}"
+        measured = {
+            q
+            for op in called_circuit.all_operations()
+            if isinstance(op.gate, cirq.MeasurementGate)
+            for q in op.qubits
+        }
+        assert measured == set(pauli_qubits)
+
+    # Now test with symmetries
+    sym_qubit = qubits[0]
+    params = CircuitToPauliStringsParameters(
+        circuit=circuit,
+        pauli_strings=[(pauli_str,)],
+        postselection_symmetries=[(cirq.PauliString(cirq.Z(sym_qubit)), 1)],
+    )
+
+    mock_sampler_sym = MagicMock()
+    mock_res_sym = ResultDict(
+        params=cirq.ParamResolver({}), measurements={"result": np.zeros((1, 4))}
+    )
+    mock_sampler_sym.run_batch.return_value = [[mock_res_sym]]
+    mock_sampler_sym.run_sweep.return_value = [mock_res_sym]
+
+    measure_pauli_strings([params], mock_sampler_sym, 10, 10, 0, 1234, use_sweep=use_sweep)
+
+    if use_sweep:
+        called_circuits_sym = [
+            call.kwargs['program'] for call in mock_sampler_sym.run_sweep.call_args_list
+        ]
+    else:
+        called_circuits_sym = [
+            c for batch in mock_sampler_sym.run_batch.call_args_list for c in batch.args[0]
+        ]
+
+    expected_qubits = set(pauli_qubits) | {sym_qubit}
+    for called_circuit in called_circuits_sym:
+        measured = {
+            q
+            for op in called_circuit.all_operations()
+            if isinstance(op.gate, cirq.MeasurementGate)
+            for q in op.qubits
+        }
+        assert measured == expected_qubits

From 3873587a4d9412a164d25b05a59d95332f0e7c91 Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Wed, 15 Apr 2026 10:30:36 -0700
Subject: [PATCH 03/11] Resolve lint (?)

---
 ...ing_measurement_with_readout_mitigation.py | 26 +++++++++++--------
 ...easurement_with_readout_mitigation_test.py | 19 ++++++--------
 2 files changed, 23 insertions(+), 22 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 a5d56c89d42..342059b04bc 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
@@ -26,7 +26,7 @@
 import sympy
 
 import cirq.contrib.shuffle_circuits.shuffle_circuits_with_readout_benchmarking as sc_readout
-from cirq import circuits, ops, work
+from cirq import circuits, ops, study, work
 from cirq.experiments.readout_confusion_matrix import TensoredConfusionMatrices
 
 if TYPE_CHECKING:
@@ -834,7 +834,7 @@ def _measure_pauli_strings_with_symmetries(
         )
 
         # Run the sweeps
-        all_circuits_measurement_results: list[list[study.Result]] = []
+        all_circuits_measurement_results: list[Sequence[study.Result]] = []
         for parameterized_circuit, sweep in zip(pauli_measurement_circuits, sweep_params):
             results_for_one_circuit = sampler.run_sweep(
                 program=parameterized_circuit, params=sweep, repetitions=pauli_repetitions
@@ -883,10 +883,12 @@ def _measure_pauli_strings_with_symmetries(
                 + len(circuit_to_pauli_params.pauli_strings)
             ]
             qubits_in_circuit = tuple(sorted(circuit_to_pauli_params.circuit.all_qubits()))
-            single_circuit_pauli_measurement_results: list[PauliStringMeasurementResult] = []
+            single_circuit_pauli_measurement_results_in_batch: list[
+                PauliStringMeasurementResult
+            ] = []
 
             for i, circuit_result in enumerate(circuit_results):
-                single_circuit_pauli_measurement_results.extend(
+                single_circuit_pauli_measurement_results_in_batch.extend(
                     _process_symmetry_measurement_results(
                         qubits_in_circuit,
                         circuit_to_pauli_params.pauli_strings[i],
@@ -900,7 +902,7 @@ def _measure_pauli_strings_with_symmetries(
             final_measurement_results.append(
                 CircuitToPauliStringsMeasurementResult(
                     circuit=circuit_to_pauli_params.circuit,
-                    results=single_circuit_pauli_measurement_results,
+                    results=single_circuit_pauli_measurement_results_in_batch,
                 )
             )
     return final_measurement_results
@@ -1039,7 +1041,9 @@ def _measure_pauli_strings_with_confusion_matrices(
 
     for circuit_to_pauli in normalized_circuits_to_pauli:
         input_circuit = circuit_to_pauli.circuit
-        pauli_string_groups = circuit_to_pauli.pauli_strings
+        pauli_string_groups_in_circuit: tuple[tuple[ops.PauliString]] = (
+            circuit_to_pauli.pauli_strings
+        )
 
         disable_readout_mitigation = num_random_bitstrings == 0
 
@@ -1051,16 +1055,16 @@ def _measure_pauli_strings_with_confusion_matrices(
                 input_circuit_index += 1
             else:
                 results_slice = slice(
-                    circuit_result_index, circuit_result_index + len(pauli_string_groups)
+                    circuit_result_index, circuit_result_index + len(pauli_string_groups_in_circuit)
                 )
                 circuits_results_for_group = [r[0] for r in sweep_circuits_results[results_slice]]
-                circuit_result_index += len(pauli_string_groups)
+                circuit_result_index += len(pauli_string_groups_in_circuit)
         else:
             results_slice = slice(
-                circuit_result_index, circuit_result_index + len(pauli_string_groups)
+                circuit_result_index, circuit_result_index + len(pauli_string_groups_in_circuit)
             )
             circuits_results_for_group = circuits_results[results_slice]
-            circuit_result_index += len(pauli_string_groups)
+            circuit_result_index += len(pauli_string_groups_in_circuit)
 
         fixed_calibration_key = (
             tuple(qubits_to_measure_arg)
@@ -1068,7 +1072,7 @@ def _measure_pauli_strings_with_confusion_matrices(
             else None
         )
         pauli_measurement_results = _process_pauli_measurement_results(
-            pauli_string_groups,
+            pauli_string_groups_in_circuit,
             circuits_results_for_group,
             calibration_results,
             pauli_repetitions,
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 846e662c842..92ecc4572d8 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
@@ -319,8 +319,8 @@ def test_group_pauli_string_measurement_errors_with_noise(use_sweep: bool) -> No
 @pytest.mark.parametrize("use_sweep", [True, False])
 def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
     """Test mixed input: some circuits using confusion matrices, some using symmetries.
-    
-    This test specifically includes a QWC group with multiple Pauli strings to 
+
+    This test specifically includes a QWC group with multiple Pauli strings to
     ensure the processing logic handles nested groups correctly.
     """
     qubits_1 = cirq.LineQubit.range(3)
@@ -329,23 +329,20 @@ def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
     # Circuit 1 is a Standard GHZ without symmetry.
     circuit_1 = cirq.FrozenCircuit(_create_ghz(3, qubits_1))
     params_1 = CircuitToPauliStringsParameters(
-        circuit=circuit_1, 
-        pauli_strings=((cirq.PauliString(cirq.Z(qubits_1[0])),),)
+        circuit=circuit_1, pauli_strings=((cirq.PauliString(cirq.Z(qubits_1[0])),),)
     )
 
     # Circuit 2 is a GHZ with a known symmetry (Z0*Z1 = 1)
     circuit_2 = cirq.FrozenCircuit(_create_ghz(5, qubits_2))
     symmetry: cirq.PauliString = cirq.PauliString(cirq.Z(qubits_2[0]) * cirq.Z(qubits_2[1]))
 
-    pauli_group = (
+    pauli_group: tuple[cirq.PauliString, ...] = (
         cirq.PauliString(cirq.Z(qubits_2[0])),
-        2.0 * cirq.PauliString(cirq.Z(qubits_2[1]))
+        2.0 * cirq.PauliString(cirq.Z(qubits_2[1])),
     )
 
     params_sym = CircuitToPauliStringsParameters(
-        circuit=circuit_2,
-        pauli_strings=(pauli_group,), 
-        postselection_symmetries=((symmetry, 1),),
+        circuit=circuit_2, pauli_strings=(pauli_group,), postselection_symmetries=((symmetry, 1),)
     )
 
     sampler = NoisySingleQubitReadoutSampler(p0=0.01, p1=0.02, seed=1234)
@@ -366,7 +363,7 @@ def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
         # For Circuit 2, we expect exactly 2 results because the group had 2 Pauli strings
         if circuit_res.circuit == circuit_2:
             assert len(circuit_res.results) == 2
-        
+
         for res in circuit_res.results:
             if circuit_res.circuit == circuit_2:
                 assert isinstance(res.calibration_result, PostFilteringSymmetryCalibrationResult)
@@ -855,7 +852,7 @@ def test_sampler_receives_correct_circuits(use_sweep: bool) -> None:
     qubits = cirq.LineQubit.range(5)
     circuit = cirq.FrozenCircuit(_create_ghz(5, qubits))
     pauli_qubits = qubits[1:4]  # Q1, Q2, Q3
-    pauli_str = cirq.PauliString(dict.fromkeys(pauli_qubits, cirq.X))
+    pauli_str: cirq.PauliString = cirq.PauliString(dict.fromkeys(pauli_qubits, cirq.X))
 
     # Test standard Pauli String without Symmetries
     circuits_to_pauli = {circuit: [pauli_str]}

From fccb2350bf219beb3fcc0602017cdb475716911d Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Wed, 15 Apr 2026 10:38:25 -0700
Subject: [PATCH 04/11] Update doc string

---
 ...tring_measurement_with_readout_mitigation.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 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 342059b04bc..9ba0abfccd4 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
@@ -40,6 +40,10 @@
 class CircuitToPauliStringsParameters:
     """Parameters for measuring Pauli strings on a circuit.
 
+    If postselection symmetries are provided, this circuit will be measured using
+    the post-selection symmetry method. If no postselection symmetries are provided,
+    this circuit will be measured using the confusion matrix method.
+
     Attributes:
         circuit: The circuit to measure.
         pauli_strings:
@@ -1107,6 +1111,10 @@ def measure_pauli_strings(
     """Measures expectation values of Pauli strings on given circuits with/without
     readout error mitigation.
 
+    Note: If `postselection_symmetries` are included in the `circuits_to_pauli` parameters,
+    readout benchmarking is not performed. In this case, the `readout_repetitions` and
+    `num_random_bitstrings` arguments are ignored.
+
     Args:
         circuits_to_pauli: A list of CircuitToPauliStringsParameters objects, where each object contains:
             - The circuit to measure.
@@ -1116,14 +1124,19 @@ def measure_pauli_strings(
         pauli_repetitions: The number of repetitions for each circuit when measuring
             Pauli strings.
         readout_repetitions: The number of repetitions for readout calibration
-            in the shuffled benchmarking.
+            in the shuffled benchmarking. (Ignored if `postselection_symmetries` are provided).
         num_random_bitstrings: The number of random bitstrings to use in readout
-            benchmarking.
+            benchmarking. (Ignored if `postselection_symmetries` are provided).
         rng_or_seed: A random number generator or seed for the readout benchmarking.
         use_sweep: If True, uses parameterized circuits and sweeps parameters
             for both Pauli measurements and readout benchmarking. Defaults to False.
         insert_strategy: The strategy for inserting measurement operations into the circuit.
             Defaults to circuits.InsertStrategy.INLINE.
+        measure_on_full_support: If True, calculates the union of all qubits used in all
+            Pauli strings (the full support). All circuits will then measure this full set
+            of qubits, and readout benchmarking will be performed only once on this full set,
+            rather than for every unique subset of Pauli qubits. This significantly reduces
+            overhead when measuring many Pauli strings with varying support.
 
     Returns:
         A list of CircuitToPauliStringsMeasurementResult objects, where each object contains:

From f060e469b5f21e95b06403fcd6df09b0006037ca Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Wed, 15 Apr 2026 10:41:22 -0700
Subject: [PATCH 05/11] update doc string

---
 .../pauli_string_measurement_with_readout_mitigation.py       | 4 ++--
 1 file changed, 2 insertions(+), 2 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 9ba0abfccd4..da07db918b7 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
@@ -1112,8 +1112,8 @@ def measure_pauli_strings(
     readout error mitigation.
 
     Note: If `postselection_symmetries` are included in the `circuits_to_pauli` parameters,
-    readout benchmarking is not performed. In this case, the `readout_repetitions` and
-    `num_random_bitstrings` arguments are ignored.
+    the circuit will be measured using the post-selection symmetry method.
+    In this case, the `readout_repetitions` and `num_random_bitstrings` arguments are ignored.
 
     Args:
         circuits_to_pauli: A list of CircuitToPauliStringsParameters objects, where each object contains:

From 8ac828388512ceb057c8c85eeeef08068637afca Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Wed, 15 Apr 2026 14:25:15 -0700
Subject: [PATCH 06/11] fix lint and coverage

---
 ...ing_measurement_with_readout_mitigation.py | 29 ++++++++------
 ...easurement_with_readout_mitigation_test.py | 40 ++++++++++++++-----
 2 files changed, 47 insertions(+), 22 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 da07db918b7..f21213ad926 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
@@ -19,7 +19,7 @@
 import itertools
 import time
 from collections.abc import Mapping, Sequence
-from typing import cast, TYPE_CHECKING, Union
+from typing import cast, TYPE_CHECKING
 
 import attrs
 import numpy as np
@@ -373,7 +373,7 @@ def _extract_readout_qubits(
     symmetries: Sequence[ops.PauliString | ops.PauliSum] | None = None,
 ) -> list[ops.Qid]:
     """Extracts unique qubits from both the target Pauli strings and the symmetries."""
-    all_qubits = set(q for ps in pauli_strings for q in ps.qubits)
+    all_qubits = {q for ps in pauli_strings for q in ps.qubits}
 
     if symmetries:
         for sym in symmetries:
@@ -383,11 +383,11 @@ def _extract_readout_qubits(
                 for term in sym:
                     all_qubits.update(term.qubits)
 
-    return sorted(list(all_qubits))
+    return sorted(all_qubits)
 
 
 def _pauli_objs_to_basis_change_ops(
-    pauli_objs: Sequence[Union[ops.PauliString, ops.PauliSum]], qid_list: Sequence[ops.Qid]
+    pauli_objs: Sequence[ops.PauliString | ops.PauliSum], qid_list: Sequence[ops.Qid]
 ):
     flattened_terms = []
     for obj in pauli_objs:
@@ -404,7 +404,7 @@ def _pauli_objs_to_basis_change_ops(
             if pauli_op == ops.X:
                 operations.append(ops.ry(-np.pi / 2)(qubit))
                 break
-            elif pauli_op == ops.Y:
+            if pauli_op == ops.Y:
                 operations.append(ops.rx(np.pi / 2)(qubit))
                 break
 
@@ -412,7 +412,7 @@ def _pauli_objs_to_basis_change_ops(
 
 
 def _pauli_objs_to_basis_change_with_sweep(
-    pauli_objs: Sequence[Union[ops.PauliString | ops.PauliSum]], qid_list: Sequence[ops.Qid]
+    pauli_objs: Sequence[ops.PauliString | ops.PauliSum], qid_list: Sequence[ops.Qid]
 ) -> dict[str, float]:
     """Decide single-qubit rotation sweep parameters for basis change.
 
@@ -440,7 +440,7 @@ def _pauli_objs_to_basis_change_with_sweep(
                 params_dict[f"phi{qid}"] = 0.0
                 params_dict[f"theta{qid}"] = 1 / 2
                 break
-            elif pauli_op == ops.Y:
+            if pauli_op == ops.Y:
                 params_dict[f"phi{qid}"] = 1.0
                 params_dict[f"theta{qid}"] = 1 / 2
                 break
@@ -601,8 +601,7 @@ def _split_input_circuits(
             # If no postselection symmetries are provided, treat the circuit as a confusion circuit
             confusion_circuits.append(circuit_to_pauli_params)
             continue
-        else:
-            symmetry_circuits.append(circuit_to_pauli_params)
+        symmetry_circuits.append(circuit_to_pauli_params)
     return symmetry_circuits, confusion_circuits
 
 
@@ -806,7 +805,8 @@ def _measure_pauli_strings_with_symmetries(
 ) -> list[CircuitToPauliStringsMeasurementResult]:
     """
     Measures expectation values of Pauli strings on given circuits with postselection symmetries.
-    This function takes a list of CircuitToPauliStringsParameters. Each parameter contains a circuit, its associated list of QWC Pauli string groups and postselection symmetries.
+    This function takes a list of CircuitToPauliStringsParameters. Each parameter contains
+    a circuit, its associated list of QWC Pauli string groups and postselection symmetries.
     For each circuit_to_pauli, it:
     1. Runs the circuits to get the measurement results.
     2. Filters the measurement results based on postselection symmetries.
@@ -1116,10 +1116,13 @@ def measure_pauli_strings(
     In this case, the `readout_repetitions` and `num_random_bitstrings` arguments are ignored.
 
     Args:
-        circuits_to_pauli: A list of CircuitToPauliStringsParameters objects, where each object contains:
+        circuits_to_pauli: A list of CircuitToPauliStringsParameters objects, where each
+        object contains:
             - The circuit to measure.
-            - A list of QWC groups (list[list[ops.PauliString]]) or a list of PauliStrings (list[ops.PauliString]).
-            - A dictionary mapping Pauli strings or Pauli sums to expected eigen value for postselection symmetries.
+            - A list of QWC groups (list[list[ops.PauliString]]) or a list of PauliStrings
+            (list[ops.PauliString]).
+            - A dictionary mapping Pauli strings or Pauli sums to expected eigen value for
+            postselection symmetries.
         sampler: The sampler to use.
         pauli_repetitions: The number of repetitions for each circuit when measuring
             Pauli strings.
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 92ecc4572d8..517f7216813 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
@@ -325,6 +325,7 @@ def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
     """
     qubits_1 = cirq.LineQubit.range(3)
     qubits_2 = cirq.LineQubit.range(5)
+    qubits_3 = cirq.LineQubit.range(2)
 
     # Circuit 1 is a Standard GHZ without symmetry.
     circuit_1 = cirq.FrozenCircuit(_create_ghz(3, qubits_1))
@@ -344,11 +345,21 @@ def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
     params_sym = CircuitToPauliStringsParameters(
         circuit=circuit_2, pauli_strings=(pauli_group,), postselection_symmetries=((symmetry, 1),)
     )
+    
+    # Circuit 3 is a |+>|+> state with a PauliSum symmetry (X0 + X1 = 2).
+    circuit_3 = cirq.FrozenCircuit(cirq.H(qubits_3[0]), cirq.H(qubits_3[1]))
+    symmetry_sum = cirq.X(qubits_3[0]) + cirq.X(qubits_3[1])
+    params_sym_sum = CircuitToPauliStringsParameters(
+        circuit=circuit_3,
+        pauli_strings=((2.0 * cirq.PauliString(cirq.X(qubits_3[0])),),),
+        postselection_symmetries=((symmetry_sum, 2),),
+    )
 
     sampler = NoisySingleQubitReadoutSampler(p0=0.01, p1=0.02, seed=1234)
+    simulator = cirq.Simulator()
 
     results = measure_pauli_strings(
-        [params_1, params_sym],
+        [params_1, params_sym, params_sym_sum],
         sampler,
         pauli_repetitions=500,
         readout_repetitions=500,
@@ -357,24 +368,35 @@ def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
         use_sweep=use_sweep,
     )
 
-    assert len(results) == 2
+    assert len(results) == 3
 
     for circuit_res in results:
         # For Circuit 2, we expect exactly 2 results because the group had 2 Pauli strings
         if circuit_res.circuit == circuit_2:
             assert len(circuit_res.results) == 2
 
+        # Simulate the ideal circuit to extract the ground truth state vector
+        expected_val_simulation = simulator.simulate(circuit_res.circuit.unfreeze())
+        final_state_vector = expected_val_simulation.final_state_vector
+
         for res in circuit_res.results:
-            if circuit_res.circuit == circuit_2:
+            # Calculate the ideal expectation directly from the state vector and Pauli string
+            ideal_expectation = _ideal_expectation_based_on_pauli_string(
+                res.pauli_string, final_state_vector
+            )
+
+            # Assert the mitigated result falls statistically close to the ideal simulation
+            assert np.isclose(
+                res.mitigated_expectation,
+                ideal_expectation,
+                atol=10 * res.mitigated_stddev,
+            )
+
+            # Maintain type validations based on whether symmetries were used
+            if circuit_res.circuit in [circuit_2, circuit_3]:
                 assert isinstance(res.calibration_result, PostFilteringSymmetryCalibrationResult)
-                # Verify coefficient was applied (2.0 for the second Pauli string)
-                if res.pauli_string.coefficient == 2.0:
-                    # Expectation should be roughly 2.0 * 1.0 or 2.0 * -1.0 depending on state
-                    # for GHZ Z expectation is 0, so 2.0 * 0 = 0.
-                    assert -2.2 <= res.mitigated_expectation <= 2.2
             else:
                 assert isinstance(res.calibration_result, SingleQubitReadoutCalibrationResult)
-                assert -1.1 <= res.mitigated_expectation <= 1.1
 
 
 @pytest.mark.parametrize("use_sweep", [True, False])

From e16729b00fc856006ae862347bd076504ae3f8cb Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Wed, 15 Apr 2026 14:35:40 -0700
Subject: [PATCH 07/11] fix format

---
 .../pauli_string_measurement_with_readout_mitigation.py     | 3 +--
 ...pauli_string_measurement_with_readout_mitigation_test.py | 6 ++----
 2 files changed, 3 insertions(+), 6 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 f21213ad926..8196d663b4e 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
@@ -803,8 +803,7 @@ def _measure_pauli_strings_with_symmetries(
     use_sweep: bool,
     insert_strategy: circuits.InsertStrategy,
 ) -> list[CircuitToPauliStringsMeasurementResult]:
-    """
-    Measures expectation values of Pauli strings on given circuits with postselection symmetries.
+    """Measures expectation values of Pauli strings on given circuits with postselection symmetries.
     This function takes a list of CircuitToPauliStringsParameters. Each parameter contains
     a circuit, its associated list of QWC Pauli string groups and postselection symmetries.
     For each circuit_to_pauli, it:
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 517f7216813..f29e6d13958 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
@@ -345,7 +345,7 @@ def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
     params_sym = CircuitToPauliStringsParameters(
         circuit=circuit_2, pauli_strings=(pauli_group,), postselection_symmetries=((symmetry, 1),)
     )
-    
+
     # Circuit 3 is a |+>|+> state with a PauliSum symmetry (X0 + X1 = 2).
     circuit_3 = cirq.FrozenCircuit(cirq.H(qubits_3[0]), cirq.H(qubits_3[1]))
     symmetry_sum = cirq.X(qubits_3[0]) + cirq.X(qubits_3[1])
@@ -387,9 +387,7 @@ def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
 
             # Assert the mitigated result falls statistically close to the ideal simulation
             assert np.isclose(
-                res.mitigated_expectation,
-                ideal_expectation,
-                atol=10 * res.mitigated_stddev,
+                res.mitigated_expectation, ideal_expectation, atol=10 * res.mitigated_stddev
             )
 
             # Maintain type validations based on whether symmetries were used

From bf0379ef0e50307d474c85894eecf74b200ba8ae Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Tue, 21 Apr 2026 16:27:16 -0700
Subject: [PATCH 08/11] resolve comments, not completed

---
 .../cirq/contrib/paulistring/__init__.py      |   2 +-
 ...ing_measurement_with_readout_mitigation.py | 259 +++++++++---------
 ...easurement_with_readout_mitigation_test.py |  10 +-
 3 files changed, 137 insertions(+), 134 deletions(-)

diff --git a/cirq-core/cirq/contrib/paulistring/__init__.py b/cirq-core/cirq/contrib/paulistring/__init__.py
index cd8481eb22b..678f6b7ca7d 100644
--- a/cirq-core/cirq/contrib/paulistring/__init__.py
+++ b/cirq-core/cirq/contrib/paulistring/__init__.py
@@ -46,5 +46,5 @@
 from cirq.contrib.paulistring.pauli_string_measurement_with_readout_mitigation import (
     measure_pauli_strings as measure_pauli_strings,
     CircuitToPauliStringsParameters as CircuitToPauliStringsParameters,
-    PostFilteringSymmetryCalibrationResult as PostFilteringSymmetryCalibrationResult,
+    CircuitToPauliStringsMeasurementResult as CircuitToPauliStringsMeasurementResult,
 )
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 8196d663b4e..64009187aa1 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
@@ -18,8 +18,8 @@
 
 import itertools
 import time
-from collections.abc import Mapping, Sequence
-from typing import cast, TYPE_CHECKING
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Iterator, cast, TYPE_CHECKING, overload
 
 import attrs
 import numpy as np
@@ -46,24 +46,34 @@ class CircuitToPauliStringsParameters:
 
     Attributes:
         circuit: The circuit to measure.
-        pauli_strings:
-            - A tuple of QWC groups (tuple[tuple[ops.PauliString]]). Each QWC group
-              is a tuple of PauliStrings that are mutually Qubit-Wise Commuting.
-              Pauli strings within the same group will be calculated using the
-              same measurement results.
-        postselection_symmetries: A tuple mapping Pauli strings or Pauli sums to
-                                  expected values for postselection symmetries.
-                                  Measured bitstrings which do not have the indicated
-                                  values of the symmetry operators are postselected out.
+        pauli_strings: The tuple of QWC groups (tuple[tuple[ops.PauliString, ...], ...]).
+            Each QWC group is a tuple of PauliStrings that are mutually Qubit-Wise Commuting.
+            Pauli strings within the same group will be calculated using the same measurement
+            results.
+        postselection_symmetries: The tuple mapping Pauli strings or Pauli sums to expected values
+            for postselection symmetries.  Measured bitstrings which do not have the indicated
+            values of the symmetry operators are postselected out.
+        readout_repetitions: The number of repetitions to use for readout calibration.
+        num_random_bitstrings: The number of random bitstrings to use for readout mitigation.
     """
 
     circuit: circuits.FrozenCircuit
-    pauli_strings: tuple[tuple[ops.PauliString]] = attrs.field(
+    pauli_strings: tuple[tuple[ops.PauliString, ...], ...] = attrs.field(
         converter=lambda groups: tuple(tuple(group) for group in groups)
     )
     postselection_symmetries: Sequence[tuple[ops.PauliString | ops.PauliSum, int]] = attrs.field(
         default=()
     )
+    readout_repetitions: int = 1
+    num_random_bitstrings: int = 0
+    
+    def __attrs_post_init__(self):
+        # Check num_random_bitstrings is bigger than or equal to 0
+        if self.num_random_bitstrings is not None and self.num_random_bitstrings < 0:
+            raise ValueError("Must provide zero or more num_random_bitstrings.")
+        # Check readout_repetitions is bigger than 0
+        if self.readout_repetitions is not None and self.readout_repetitions <= 0:
+            raise ValueError("Must provide positive readout_repetitions for readout calibration.")
 
 
 @attrs.frozen
@@ -118,6 +128,16 @@ class CircuitToPauliStringsMeasurementResult:
     results: list[PauliStringMeasurementResult]
 
 
+def _flatten_pauli_objs(
+    pauli_objs: Sequence[ops.PauliString | ops.PauliSum],
+) -> Iterator[ops.PauliString]:
+    for obj in pauli_objs:
+        if isinstance(obj, ops.PauliString):
+            yield obj
+        else:
+            yield from obj
+
+
 def _commute_or_identity(
     op1: ops.Pauli | ops.IdentityGate, op2: ops.Pauli | ops.IdentityGate
 ) -> bool:
@@ -286,14 +306,14 @@ def _validate_input(
         | list[CircuitToPauliStringsParameters]
     ),
     pauli_repetitions: int,
-    readout_repetitions: int,
-    num_random_bitstrings: int,
     rng_or_seed: np.random.Generator | int,
+    readout_repetitions: int | None = None,
+    num_random_bitstrings: int | None = None,
 ):
     if not circuits_to_pauli:
         raise ValueError("Input circuits_to_pauli parameter must not be empty.")
 
-    normalized_circuits_to_pauli = _validate_and_normalize_unformatted_input(circuits_to_pauli)
+    normalized_circuits_to_pauli = _validate_and_normalize_unformatted_input(circuits_to_pauli, readout_repetitions, num_random_bitstrings)
 
     _validate_circuit_to_pauli_strings_parameters(normalized_circuits_to_pauli)
 
@@ -305,14 +325,6 @@ def _validate_input(
     if pauli_repetitions <= 0:
         raise ValueError("Must provide positive pauli_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 positive readout_repetitions for readout calibration.")
-
     return normalized_circuits_to_pauli
 
 
@@ -322,6 +334,8 @@ def _validate_and_normalize_unformatted_input(
         | Mapping[circuits.FrozenCircuit, Sequence[Sequence[ops.PauliString]]]
         | list[CircuitToPauliStringsParameters]
     ),
+    readout_repetitions: int | None = None,
+    num_random_bitstrings: int | None = None,
 ) -> list[CircuitToPauliStringsParameters]:
     """Converts any valid input format into a standardized list of parameters
     where pauli_strings is always Sequence[Sequence[PauliString]]."""
@@ -342,6 +356,8 @@ def _validate_and_normalize_unformatted_input(
                     circuit=circuit,
                     pauli_strings=cast(tuple[tuple[ops.PauliString]], normalized_paulis),
                     postselection_symmetries=[],
+                    readout_repetitions=readout_repetitions,
+                    num_random_bitstrings=num_random_bitstrings,
                 )
             )
     elif isinstance(circuits_input, list):
@@ -370,18 +386,14 @@ def _validate_and_normalize_unformatted_input(
 
 def _extract_readout_qubits(
     pauli_strings: Sequence[ops.PauliString],
-    symmetries: Sequence[ops.PauliString | ops.PauliSum] | None = None,
+    symmetries: Sequence[ops.PauliString | ops.PauliSum] = (),
 ) -> list[ops.Qid]:
     """Extracts unique qubits from both the target Pauli strings and the symmetries."""
     all_qubits = {q for ps in pauli_strings for q in ps.qubits}
 
     if symmetries:
-        for sym in symmetries:
-            if isinstance(sym, ops.PauliString):
-                all_qubits.update(sym.qubits)
-            elif isinstance(sym, ops.PauliSum):
-                for term in sym:
-                    all_qubits.update(term.qubits)
+        for sym in _flatten_pauli_objs(symmetries):
+            all_qubits.update(sym.qubits)
 
     return sorted(all_qubits)
 
@@ -389,12 +401,7 @@ def _extract_readout_qubits(
 def _pauli_objs_to_basis_change_ops(
     pauli_objs: Sequence[ops.PauliString | ops.PauliSum], qid_list: Sequence[ops.Qid]
 ):
-    flattened_terms = []
-    for obj in pauli_objs:
-        if isinstance(obj, ops.PauliString):
-            flattened_terms.append(obj)
-        elif isinstance(obj, ops.PauliSum):
-            flattened_terms.extend(list(obj))
+    flattened_terms = list(_flatten_pauli_objs(pauli_objs))
 
     operations = []
 
@@ -422,12 +429,7 @@ def _pauli_objs_to_basis_change_with_sweep(
     Returns:
         A dictionary mapping parameter names to their values for basis change.
     """
-    flattened_terms = []
-    for obj in pauli_objs:
-        if isinstance(obj, ops.PauliString):
-            flattened_terms.append(obj)
-        elif isinstance(obj, ops.PauliSum):
-            flattened_terms.extend(list(obj))
+    flattened_terms = list(_flatten_pauli_objs(pauli_objs))
 
     params_dict: dict[str, float] = {}
 
@@ -597,11 +599,10 @@ def _split_input_circuits(
     confusion_circuits: list[CircuitToPauliStringsParameters] = []
 
     for circuit_to_pauli_params in circuits_to_pauli_params:
-        if not circuit_to_pauli_params.postselection_symmetries:
-            # If no postselection symmetries are provided, treat the circuit as a confusion circuit
-            confusion_circuits.append(circuit_to_pauli_params)
-            continue
-        symmetry_circuits.append(circuit_to_pauli_params)
+        if circuit_to_pauli_params.postselection_symmetries:
+            symmetry_circuits.append(circuit_to_pauli_params)
+        else:
+            confusion_circuits.append(circuit_to_pauli_params) 
     return symmetry_circuits, confusion_circuits
 
 
@@ -615,34 +616,29 @@ def _process_symmetry_measurement_results(
     """Filters measurement results using symmetries and calculates expectations."""
     single_circuit_pauli_measurement_results: list[PauliStringMeasurementResult] = []
 
+    qubit_to_index = {q: i for i, q in enumerate(qubits)}
+
+    def _qubits_to_indices(qobjs: Iterable[ops.Qid]) -> list[int]:
+        return [qubit_to_index[q] for q in qobjs]
+
     # filter out bitstrings based on postselection symmetries
     measurement_result_eigenvalues = 1 - 2 * measurement_results
     rows_to_keep_mask = np.ones(len(measurement_result_eigenvalues), dtype=bool)
 
     for sym, expected_value in circuit_to_pauli.postselection_symmetries:
-        # Determine which rows to keep based on the symmetry
-        if isinstance(sym, ops.PauliString):
-            sym_qubit_indices = [qubits.index(q) for q in sym.keys()]
-            actual_eigenvalues = np.prod(
-                measurement_result_eigenvalues[:, sym_qubit_indices], axis=1
+        sum_eigenvalues = sum(
+            ps.coefficient.real * np.prod(
+                measurement_result_eigenvalues[:, _qubits_to_indices(ps.keys())], axis=1
             )
-            rows_to_keep_mask &= actual_eigenvalues == expected_value
-
-        elif isinstance(sym, ops.PauliSum):
-            sum_eigenvalues = np.zeros(len(measurement_result_eigenvalues), dtype=float)
-            for term in sym:
-                term_qubit_indices = [qubits.index(q) for q in term.keys()]
-                term_eigenvalues = np.prod(
-                    measurement_result_eigenvalues[:, term_qubit_indices], axis=1
-                )
-                sum_eigenvalues += term.coefficient.real * term_eigenvalues
-            rows_to_keep_mask &= np.isclose(sum_eigenvalues, expected_value)
+            for ps in _flatten_pauli_objs([sym])
+        )
+        rows_to_keep_mask &= np.isclose(sum_eigenvalues, expected_value)
 
     post_selection_circuits_results = measurement_results[rows_to_keep_mask]
 
     for pauli_str in pauli_string_groups:
         qubits_sorted = sorted(pauli_str.qubits)
-        qubit_indices = [qubits.index(q) for q in qubits_sorted]
+        qubit_indices = _qubits_to_indices(qubits_sorted)
         relevant_bits_unmit = measurement_results[:, qubit_indices]
 
         if len(post_selection_circuits_results) == 0:
@@ -729,6 +725,8 @@ def _process_pauli_measurement_results(
         else:
             pauli_readout_qubits = _extract_readout_qubits(pauli_strs)
             calibration_key = tuple(pauli_readout_qubits)
+            
+        qubit_to_index = {q: i for i, q in enumerate(pauli_readout_qubits)}
 
         calibration_result = (
             calibration_results[calibration_key] if not disable_readout_mitigation else None
@@ -736,7 +734,7 @@ def _process_pauli_measurement_results(
 
         for pauli_str in pauli_strs:
             qubits_sorted = sorted(pauli_str.qubits)
-            qubit_indices = [pauli_readout_qubits.index(q) for q in qubits_sorted]
+            qubit_indices = [qubit_to_index[q] for q in qubits_sorted]
 
             if disable_readout_mitigation:
                 pauli_str_calibration_result = None
@@ -828,8 +826,7 @@ def _measure_pauli_strings_with_symmetries(
     if not circuits_to_pauli:
         return []
 
-    final_measurement_results: list[CircuitToPauliStringsMeasurementResult] = []
-
+    all_circuits_measurement_results: list[Sequence[study.Result]] = []
     # Generate measurement circuits
     if use_sweep:
         pauli_measurement_circuits, sweep_params = _generate_basis_change_circuits_with_sweep(
@@ -837,39 +834,13 @@ def _measure_pauli_strings_with_symmetries(
         )
 
         # Run the sweeps
-        all_circuits_measurement_results: list[Sequence[study.Result]] = []
         for parameterized_circuit, sweep in zip(pauli_measurement_circuits, sweep_params):
             results_for_one_circuit = sampler.run_sweep(
                 program=parameterized_circuit, params=sweep, repetitions=pauli_repetitions
             )
             all_circuits_measurement_results.append(results_for_one_circuit)
-
-        # Process the results
-        for circuit_to_pauli_params, circuit_results in zip(
-            circuits_to_pauli, all_circuits_measurement_results
-        ):
-            qubits_in_circuit = tuple(sorted(circuit_to_pauli_params.circuit.all_qubits()))
-            single_circuit_pauli_measurement_results: list[PauliStringMeasurementResult] = []
-
-            for i, circuit_result in enumerate(circuit_results):
-                single_circuit_pauli_measurement_results.extend(
-                    _process_symmetry_measurement_results(
-                        qubits_in_circuit,
-                        circuit_to_pauli_params.pauli_strings[i],
-                        circuit_result.measurements["result"],
-                        circuit_to_pauli_params,
-                        pauli_repetitions,
-                    )
-                )
-
-            final_measurement_results.append(
-                CircuitToPauliStringsMeasurementResult(
-                    circuit=circuit_to_pauli_params.circuit,
-                    results=single_circuit_pauli_measurement_results,
-                )
-            )
     else:
-        # Process in batch mode
+        # Run in batch mode
         pauli_measurement_circuits = _generate_basis_change_circuits(
             circuits_to_pauli, insert_strategy
         )
@@ -881,33 +852,38 @@ def _measure_pauli_strings_with_symmetries(
 
         circuit_result_index = 0
         for circuit_to_pauli_params in circuits_to_pauli:
+            num_results = len(circuit_to_pauli_params.pauli_strings)
             circuit_results = circuits_measurement_results[
                 circuit_result_index : circuit_result_index
                 + len(circuit_to_pauli_params.pauli_strings)
             ]
-            qubits_in_circuit = tuple(sorted(circuit_to_pauli_params.circuit.all_qubits()))
-            single_circuit_pauli_measurement_results_in_batch: list[
-                PauliStringMeasurementResult
-            ] = []
-
-            for i, circuit_result in enumerate(circuit_results):
-                single_circuit_pauli_measurement_results_in_batch.extend(
-                    _process_symmetry_measurement_results(
-                        qubits_in_circuit,
-                        circuit_to_pauli_params.pauli_strings[i],
-                        circuit_result.measurements["result"],
-                        circuit_to_pauli_params,
-                        pauli_repetitions,
-                    )
-                )
-
-            circuit_result_index += len(circuit_to_pauli_params.pauli_strings)
-            final_measurement_results.append(
-                CircuitToPauliStringsMeasurementResult(
-                    circuit=circuit_to_pauli_params.circuit,
-                    results=single_circuit_pauli_measurement_results_in_batch,
+            all_circuits_measurement_results.append(circuit_results)
+            circuit_result_index += num_results
+    
+    final_measurement_results: list[CircuitToPauliStringsMeasurementResult] = []
+    for circuit_to_pauli_params, circuit_results in zip(
+        circuits_to_pauli, all_circuits_measurement_results
+    ):
+        qubits_in_circuit = tuple(sorted(circuit_to_pauli_params.circuit.all_qubits()))
+        single_circuit_pauli_measurement_results: list[PauliStringMeasurementResult] = []
+        
+        for i, circuit_result in enumerate(circuit_results):
+            single_circuit_pauli_measurement_results.extend(
+                _process_symmetry_measurement_results(
+                    qubits_in_circuit,
+                    circuit_to_pauli_params.pauli_strings[i],
+                    circuit_result.measurements["result"],
+                    circuit_to_pauli_params,
+                    pauli_repetitions,
                 )
             )
+        final_measurement_results.append(
+            CircuitToPauliStringsMeasurementResult(
+                circuit=circuit_to_pauli_params.circuit,
+                results=single_circuit_pauli_measurement_results,
+            )
+        )
+        
     return final_measurement_results
 
 
@@ -915,8 +891,6 @@ def _measure_pauli_strings_with_confusion_matrices(
     normalized_circuits_to_pauli: list[CircuitToPauliStringsParameters],
     sampler: work.Sampler,
     pauli_repetitions: int,
-    readout_repetitions: int,
-    num_random_bitstrings: int,
     rng_or_seed: np.random.Generator | int,
     use_sweep: bool = False,
     insert_strategy: circuits.InsertStrategy = circuits.InsertStrategy.INLINE,
@@ -949,10 +923,6 @@ def _measure_pauli_strings_with_confusion_matrices(
         sampler: The sampler to use.
         pauli_repetitions: The number of repetitions for each circuit when measuring
             Pauli strings.
-        readout_repetitions: The number of repetitions for readout calibration
-            in the shuffled benchmarking.
-        num_random_bitstrings: The number of random bitstrings to use in readout
-            benchmarking.
         rng_or_seed: A random number generator or seed for the readout benchmarking.
         use_sweep: If True, uses parameterized circuits and sweeps parameters
             for both Pauli measurements and readout benchmarking. Defaults to False.
@@ -971,6 +941,9 @@ def _measure_pauli_strings_with_confusion_matrices(
     """
     if not normalized_circuits_to_pauli:
         return []
+    
+    readout_repetitions = normalized_circuits_to_pauli[0].readout_repetitions
+    num_random_bitstrings = normalized_circuits_to_pauli[0].num_random_bitstrings
 
     # Extract unique qubit tuples from input pauli strings
     if measure_on_full_support:
@@ -1092,12 +1065,12 @@ def _measure_pauli_strings_with_confusion_matrices(
     return results
 
 
+@overload
 def measure_pauli_strings(
-    circuits_to_pauli: (
-        Mapping[circuits.FrozenCircuit, Sequence[ops.PauliString]]
-        | Mapping[circuits.FrozenCircuit, Sequence[Sequence[ops.PauliString]]]
-        | list[CircuitToPauliStringsParameters]
-    ),
+    circuits_to_pauli: Mapping[
+        circuits.FrozenCircuit,
+        Sequence[ops.PauliString] | Sequence[Sequence[ops.PauliString]]
+    ],
     sampler: work.Sampler,
     pauli_repetitions: int,
     readout_repetitions: int,
@@ -1106,6 +1079,36 @@ def measure_pauli_strings(
     use_sweep: bool = False,
     insert_strategy: circuits.InsertStrategy = circuits.InsertStrategy.INLINE,
     measure_on_full_support: bool = False,
+) -> list[CircuitToPauliStringsMeasurementResult]:
+    ...
+
+
+@overload
+def measure_pauli_strings(
+    circuits_to_pauli: list[CircuitToPauliStringsParameters],
+    sampler: work.Sampler,
+    pauli_repetitions: int,
+    rng_or_seed: np.random.Generator | int,
+    use_sweep: bool = False,
+    insert_strategy: circuits.InsertStrategy = circuits.InsertStrategy.INLINE,
+    measure_on_full_support: bool = False,
+) -> list[CircuitToPauliStringsMeasurementResult]:
+    ...
+    
+
+def measure_pauli_strings(
+    circuits_to_pauli: (
+        Mapping[circuits.FrozenCircuit, Sequence[ops.PauliString] | Sequence[Sequence[ops.PauliString]]]
+        | list[CircuitToPauliStringsParameters]
+    ),
+    sampler: work.Sampler,
+    pauli_repetitions: int,
+    readout_repetitions: int | None = None,
+    num_random_bitstrings: int | None = None,
+    rng_or_seed: np.random.Generator | int = 0 ,
+    use_sweep: bool = False,
+    insert_strategy: circuits.InsertStrategy = circuits.InsertStrategy.INLINE,
+    measure_on_full_support: bool = False,
 ) -> list[CircuitToPauliStringsMeasurementResult]:
     """Measures expectation values of Pauli strings on given circuits with/without
     readout error mitigation.
@@ -1150,9 +1153,9 @@ def measure_pauli_strings(
     normalized_circuits_to_pauli = _validate_input(
         circuits_to_pauli,
         pauli_repetitions,
+        rng_or_seed,
         readout_repetitions,
         num_random_bitstrings,
-        rng_or_seed,
     )
 
     # Split the input circuits into two lists based on the way they are measured.
@@ -1168,8 +1171,6 @@ def measure_pauli_strings(
         sampler=sampler,
         normalized_circuits_to_pauli=confusion_circuits,
         pauli_repetitions=pauli_repetitions,
-        readout_repetitions=readout_repetitions,
-        num_random_bitstrings=num_random_bitstrings,
         rng_or_seed=rng_or_seed,
         use_sweep=use_sweep,
         insert_strategy=insert_strategy,
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 f29e6d13958..c8832d942f1 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
@@ -25,7 +25,9 @@
 from cirq.contrib.paulistring import (
     CircuitToPauliStringsParameters,
     measure_pauli_strings,
-    PostFilteringSymmetryCalibrationResult,
+)
+from cirq.contrib.paulistring.pauli_string_measurement_with_readout_mitigation import (
+    PostFilteringSymmetryCalibrationResult as PostFilteringSymmetryCalibrationResult,
 )
 from cirq.experiments import SingleQubitReadoutCalibrationResult
 from cirq.experiments.single_qubit_readout_calibration_test import NoisySingleQubitReadoutSampler
@@ -335,15 +337,15 @@ def test_many_circuits_mixed_mitigation_types(use_sweep: bool) -> None:
 
     # Circuit 2 is a GHZ with a known symmetry (Z0*Z1 = 1)
     circuit_2 = cirq.FrozenCircuit(_create_ghz(5, qubits_2))
-    symmetry: cirq.PauliString = cirq.PauliString(cirq.Z(qubits_2[0]) * cirq.Z(qubits_2[1]))
+    symmetry: cirq.PauliString = -1.0 * cirq.PauliString(cirq.Z(qubits_2[0]) * cirq.Z(qubits_2[1]))
 
     pauli_group: tuple[cirq.PauliString, ...] = (
-        cirq.PauliString(cirq.Z(qubits_2[0])),
+        cirq.PauliString(cirq.Z(qubits_2[0]) * cirq.Z(qubits_2[1])),
         2.0 * cirq.PauliString(cirq.Z(qubits_2[1])),
     )
 
     params_sym = CircuitToPauliStringsParameters(
-        circuit=circuit_2, pauli_strings=(pauli_group,), postselection_symmetries=((symmetry, 1),)
+        circuit=circuit_2, pauli_strings=(pauli_group,), postselection_symmetries=((symmetry, -1),)
     )
 
     # Circuit 3 is a |+>|+> state with a PauliSum symmetry (X0 + X1 = 2).

From e8934f0d13dd728bef884c797b24ffc76f6caa23 Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Sun, 17 May 2026 17:08:33 -0700
Subject: [PATCH 09/11] fix test

---
 ...ing_measurement_with_readout_mitigation.py | 85 ++++++-------------
 1 file changed, 26 insertions(+), 59 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 64009187aa1..410fdc15b05 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
@@ -19,7 +19,7 @@
 import itertools
 import time
 from collections.abc import Iterable, Mapping, Sequence
-from typing import Iterator, cast, TYPE_CHECKING, overload
+from typing import Iterator, cast, TYPE_CHECKING
 
 import attrs
 import numpy as np
@@ -53,8 +53,6 @@ class CircuitToPauliStringsParameters:
         postselection_symmetries: The tuple mapping Pauli strings or Pauli sums to expected values
             for postselection symmetries.  Measured bitstrings which do not have the indicated
             values of the symmetry operators are postselected out.
-        readout_repetitions: The number of repetitions to use for readout calibration.
-        num_random_bitstrings: The number of random bitstrings to use for readout mitigation.
     """
 
     circuit: circuits.FrozenCircuit
@@ -64,16 +62,6 @@ class CircuitToPauliStringsParameters:
     postselection_symmetries: Sequence[tuple[ops.PauliString | ops.PauliSum, int]] = attrs.field(
         default=()
     )
-    readout_repetitions: int = 1
-    num_random_bitstrings: int = 0
-    
-    def __attrs_post_init__(self):
-        # Check num_random_bitstrings is bigger than or equal to 0
-        if self.num_random_bitstrings is not None and self.num_random_bitstrings < 0:
-            raise ValueError("Must provide zero or more num_random_bitstrings.")
-        # Check readout_repetitions is bigger than 0
-        if self.readout_repetitions is not None and self.readout_repetitions <= 0:
-            raise ValueError("Must provide positive readout_repetitions for readout calibration.")
 
 
 @attrs.frozen
@@ -306,14 +294,14 @@ def _validate_input(
         | list[CircuitToPauliStringsParameters]
     ),
     pauli_repetitions: int,
+    readout_repetitions: int,
+    num_random_bitstrings: int,
     rng_or_seed: np.random.Generator | int,
-    readout_repetitions: int | None = None,
-    num_random_bitstrings: int | None = None,
 ):
     if not circuits_to_pauli:
         raise ValueError("Input circuits_to_pauli parameter must not be empty.")
 
-    normalized_circuits_to_pauli = _validate_and_normalize_unformatted_input(circuits_to_pauli, readout_repetitions, num_random_bitstrings)
+    normalized_circuits_to_pauli = _validate_and_normalize_unformatted_input(circuits_to_pauli)
 
     _validate_circuit_to_pauli_strings_parameters(normalized_circuits_to_pauli)
 
@@ -325,6 +313,14 @@ def _validate_input(
     if pauli_repetitions <= 0:
         raise ValueError("Must provide positive pauli_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 positive readout_repetitions for readout calibration.")
+
     return normalized_circuits_to_pauli
 
 
@@ -334,8 +330,6 @@ def _validate_and_normalize_unformatted_input(
         | Mapping[circuits.FrozenCircuit, Sequence[Sequence[ops.PauliString]]]
         | list[CircuitToPauliStringsParameters]
     ),
-    readout_repetitions: int | None = None,
-    num_random_bitstrings: int | None = None,
 ) -> list[CircuitToPauliStringsParameters]:
     """Converts any valid input format into a standardized list of parameters
     where pauli_strings is always Sequence[Sequence[PauliString]]."""
@@ -356,8 +350,6 @@ def _validate_and_normalize_unformatted_input(
                     circuit=circuit,
                     pauli_strings=cast(tuple[tuple[ops.PauliString]], normalized_paulis),
                     postselection_symmetries=[],
-                    readout_repetitions=readout_repetitions,
-                    num_random_bitstrings=num_random_bitstrings,
                 )
             )
     elif isinstance(circuits_input, list):
@@ -891,6 +883,8 @@ def _measure_pauli_strings_with_confusion_matrices(
     normalized_circuits_to_pauli: list[CircuitToPauliStringsParameters],
     sampler: work.Sampler,
     pauli_repetitions: int,
+    readout_repetitions: int,
+    num_random_bitstrings: int,
     rng_or_seed: np.random.Generator | int,
     use_sweep: bool = False,
     insert_strategy: circuits.InsertStrategy = circuits.InsertStrategy.INLINE,
@@ -923,6 +917,10 @@ def _measure_pauli_strings_with_confusion_matrices(
         sampler: The sampler to use.
         pauli_repetitions: The number of repetitions for each circuit when measuring
             Pauli strings.
+        readout_repetitions: The number of repetitions for readout calibration
+            in the shuffled benchmarking.
+        num_random_bitstrings: The number of random bitstrings to use in readout
+            benchmarking.
         rng_or_seed: A random number generator or seed for the readout benchmarking.
         use_sweep: If True, uses parameterized circuits and sweeps parameters
             for both Pauli measurements and readout benchmarking. Defaults to False.
@@ -941,9 +939,6 @@ def _measure_pauli_strings_with_confusion_matrices(
     """
     if not normalized_circuits_to_pauli:
         return []
-    
-    readout_repetitions = normalized_circuits_to_pauli[0].readout_repetitions
-    num_random_bitstrings = normalized_circuits_to_pauli[0].num_random_bitstrings
 
     # Extract unique qubit tuples from input pauli strings
     if measure_on_full_support:
@@ -1065,47 +1060,17 @@ def _measure_pauli_strings_with_confusion_matrices(
     return results
 
 
-@overload
-def measure_pauli_strings(
-    circuits_to_pauli: Mapping[
-        circuits.FrozenCircuit,
-        Sequence[ops.PauliString] | Sequence[Sequence[ops.PauliString]]
-    ],
-    sampler: work.Sampler,
-    pauli_repetitions: int,
-    readout_repetitions: int,
-    num_random_bitstrings: int,
-    rng_or_seed: np.random.Generator | int,
-    use_sweep: bool = False,
-    insert_strategy: circuits.InsertStrategy = circuits.InsertStrategy.INLINE,
-    measure_on_full_support: bool = False,
-) -> list[CircuitToPauliStringsMeasurementResult]:
-    ...
-
-
-@overload
-def measure_pauli_strings(
-    circuits_to_pauli: list[CircuitToPauliStringsParameters],
-    sampler: work.Sampler,
-    pauli_repetitions: int,
-    rng_or_seed: np.random.Generator | int,
-    use_sweep: bool = False,
-    insert_strategy: circuits.InsertStrategy = circuits.InsertStrategy.INLINE,
-    measure_on_full_support: bool = False,
-) -> list[CircuitToPauliStringsMeasurementResult]:
-    ...
-    
-
 def measure_pauli_strings(
     circuits_to_pauli: (
-        Mapping[circuits.FrozenCircuit, Sequence[ops.PauliString] | Sequence[Sequence[ops.PauliString]]]
+        Mapping[circuits.FrozenCircuit, Sequence[ops.PauliString] | Sequence[Sequence[ops.PauliString]]
+]
         | list[CircuitToPauliStringsParameters]
     ),
     sampler: work.Sampler,
     pauli_repetitions: int,
-    readout_repetitions: int | None = None,
-    num_random_bitstrings: int | None = None,
-    rng_or_seed: np.random.Generator | int = 0 ,
+    readout_repetitions: int,
+    num_random_bitstrings: int,
+    rng_or_seed: np.random.Generator | int,
     use_sweep: bool = False,
     insert_strategy: circuits.InsertStrategy = circuits.InsertStrategy.INLINE,
     measure_on_full_support: bool = False,
@@ -1153,9 +1118,9 @@ def measure_pauli_strings(
     normalized_circuits_to_pauli = _validate_input(
         circuits_to_pauli,
         pauli_repetitions,
-        rng_or_seed,
         readout_repetitions,
         num_random_bitstrings,
+        rng_or_seed,
     )
 
     # Split the input circuits into two lists based on the way they are measured.
@@ -1171,6 +1136,8 @@ def measure_pauli_strings(
         sampler=sampler,
         normalized_circuits_to_pauli=confusion_circuits,
         pauli_repetitions=pauli_repetitions,
+        readout_repetitions=readout_repetitions,
+        num_random_bitstrings=num_random_bitstrings,
         rng_or_seed=rng_or_seed,
         use_sweep=use_sweep,
         insert_strategy=insert_strategy,

From 34fae1bdedd4368eccef3d6951695e8af228ad51 Mon Sep 17 00:00:00 2001
From: ddddddanni <danniwang@google.com>
Date: Sun, 17 May 2026 17:44:13 -0700
Subject: [PATCH 10/11] fix lint

---
 ...ing_measurement_with_readout_mitigation.py | 43 ++++++++++---------
 ...easurement_with_readout_mitigation_test.py |  5 +--
 2 files changed, 24 insertions(+), 24 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 410fdc15b05..adcdc97bf77 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
@@ -19,7 +19,8 @@
 import itertools
 import time
 from collections.abc import Iterable, Mapping, Sequence
-from typing import Iterator, cast, TYPE_CHECKING
+from typing import cast, TYPE_CHECKING
+from collections.abc import Iterator
 
 import attrs
 import numpy as np
@@ -289,8 +290,9 @@ def _validate_circuit_to_pauli_strings_parameters(
 
 def _validate_input(
     circuits_to_pauli: (
-        Mapping[circuits.FrozenCircuit, Sequence[ops.PauliString]]
-        | Mapping[circuits.FrozenCircuit, Sequence[Sequence[ops.PauliString]]]
+        Mapping[
+            circuits.FrozenCircuit, Sequence[ops.PauliString] | Sequence[Sequence[ops.PauliString]]
+        ]
         | list[CircuitToPauliStringsParameters]
     ),
     pauli_repetitions: int,
@@ -326,8 +328,9 @@ def _validate_input(
 
 def _validate_and_normalize_unformatted_input(
     circuits_input: (
-        Mapping[circuits.FrozenCircuit, Sequence[ops.PauliString]]
-        | Mapping[circuits.FrozenCircuit, Sequence[Sequence[ops.PauliString]]]
+        Mapping[
+            circuits.FrozenCircuit, Sequence[ops.PauliString] | Sequence[Sequence[ops.PauliString]]
+        ]
         | list[CircuitToPauliStringsParameters]
     ),
 ) -> list[CircuitToPauliStringsParameters]:
@@ -594,7 +597,7 @@ def _split_input_circuits(
         if circuit_to_pauli_params.postselection_symmetries:
             symmetry_circuits.append(circuit_to_pauli_params)
         else:
-            confusion_circuits.append(circuit_to_pauli_params) 
+            confusion_circuits.append(circuit_to_pauli_params)
     return symmetry_circuits, confusion_circuits
 
 
@@ -619,9 +622,8 @@ def _qubits_to_indices(qobjs: Iterable[ops.Qid]) -> list[int]:
 
     for sym, expected_value in circuit_to_pauli.postselection_symmetries:
         sum_eigenvalues = sum(
-            ps.coefficient.real * np.prod(
-                measurement_result_eigenvalues[:, _qubits_to_indices(ps.keys())], axis=1
-            )
+            ps.coefficient.real
+            * np.prod(measurement_result_eigenvalues[:, _qubits_to_indices(ps.keys())], axis=1)
             for ps in _flatten_pauli_objs([sym])
         )
         rows_to_keep_mask &= np.isclose(sum_eigenvalues, expected_value)
@@ -717,7 +719,7 @@ def _process_pauli_measurement_results(
         else:
             pauli_readout_qubits = _extract_readout_qubits(pauli_strs)
             calibration_key = tuple(pauli_readout_qubits)
-            
+
         qubit_to_index = {q: i for i, q in enumerate(pauli_readout_qubits)}
 
         calibration_result = (
@@ -845,21 +847,21 @@ def _measure_pauli_strings_with_symmetries(
         circuit_result_index = 0
         for circuit_to_pauli_params in circuits_to_pauli:
             num_results = len(circuit_to_pauli_params.pauli_strings)
-            circuit_results = circuits_measurement_results[
+            cur_circuit_results = circuits_measurement_results[
                 circuit_result_index : circuit_result_index
                 + len(circuit_to_pauli_params.pauli_strings)
             ]
-            all_circuits_measurement_results.append(circuit_results)
+            all_circuits_measurement_results.append(cur_circuit_results)
             circuit_result_index += num_results
-    
+
     final_measurement_results: list[CircuitToPauliStringsMeasurementResult] = []
-    for circuit_to_pauli_params, circuit_results in zip(
+    for circuit_to_pauli_params, circuit_results_seq in zip(
         circuits_to_pauli, all_circuits_measurement_results
     ):
         qubits_in_circuit = tuple(sorted(circuit_to_pauli_params.circuit.all_qubits()))
         single_circuit_pauli_measurement_results: list[PauliStringMeasurementResult] = []
-        
-        for i, circuit_result in enumerate(circuit_results):
+
+        for i, circuit_result in enumerate(circuit_results_seq):
             single_circuit_pauli_measurement_results.extend(
                 _process_symmetry_measurement_results(
                     qubits_in_circuit,
@@ -875,7 +877,7 @@ def _measure_pauli_strings_with_symmetries(
                 results=single_circuit_pauli_measurement_results,
             )
         )
-        
+
     return final_measurement_results
 
 
@@ -1012,7 +1014,7 @@ def _measure_pauli_strings_with_confusion_matrices(
 
     for circuit_to_pauli in normalized_circuits_to_pauli:
         input_circuit = circuit_to_pauli.circuit
-        pauli_string_groups_in_circuit: tuple[tuple[ops.PauliString]] = (
+        pauli_string_groups_in_circuit: tuple[tuple[ops.PauliString, ...], ...] = (
             circuit_to_pauli.pauli_strings
         )
 
@@ -1062,8 +1064,9 @@ def _measure_pauli_strings_with_confusion_matrices(
 
 def measure_pauli_strings(
     circuits_to_pauli: (
-        Mapping[circuits.FrozenCircuit, Sequence[ops.PauliString] | Sequence[Sequence[ops.PauliString]]
-]
+        Mapping[
+            circuits.FrozenCircuit, Sequence[ops.PauliString] | Sequence[Sequence[ops.PauliString]]
+        ]
         | list[CircuitToPauliStringsParameters]
     ),
     sampler: work.Sampler,
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 c8832d942f1..020ce616aef 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
@@ -22,10 +22,7 @@
 import pytest
 
 import cirq
-from cirq.contrib.paulistring import (
-    CircuitToPauliStringsParameters,
-    measure_pauli_strings,
-)
+from cirq.contrib.paulistring import CircuitToPauliStringsParameters, measure_pauli_strings
 from cirq.contrib.paulistring.pauli_string_measurement_with_readout_mitigation import (
     PostFilteringSymmetryCalibrationResult as PostFilteringSymmetryCalibrationResult,
 )

From 7413062410051e1a9f181f1a542941036e168063 Mon Sep 17 00:00:00 2001
From: Pavol Juhas <juhas@google.com>
Date: Mon, 18 May 2026 17:53:04 -0700
Subject: [PATCH 11/11] Format code and drop unnecessary if-statement

---
 .../pauli_string_measurement_with_readout_mitigation.py   | 8 ++------
 1 file changed, 2 insertions(+), 6 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 adcdc97bf77..c469624aedf 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
@@ -18,9 +18,8 @@
 
 import itertools
 import time
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Iterable, Iterator, Mapping, Sequence
 from typing import cast, TYPE_CHECKING
-from collections.abc import Iterator
 
 import attrs
 import numpy as np
@@ -385,10 +384,7 @@ def _extract_readout_qubits(
 ) -> list[ops.Qid]:
     """Extracts unique qubits from both the target Pauli strings and the symmetries."""
     all_qubits = {q for ps in pauli_strings for q in ps.qubits}
-
-    if symmetries:
-        for sym in _flatten_pauli_objs(symmetries):
-            all_qubits.update(sym.qubits)
+    all_qubits.update(*(sym.qubits for sym in _flatten_pauli_objs(symmetries)))
 
     return sorted(all_qubits)