Add measure_on_full_support to optimize Pauli string readout mitigation (#7760)

This PR introduces a new argument, measure_on_full_support, to the
measure_pauli_strings function. This feature significantly reduces the
overhead of readout error calibration when measuring a large number of
Pauli strings that act on different subsets of qubits.

---------

Co-authored-by: eliottrosenberg <61400172+eliottrosenberg@users.noreply.github.com>

Author: ddddddanni

cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation.py

@@ -243,7 +243,7 @@ def _pauli_strings_to_basis_change_ops(
 
 
 def _pauli_strings_to_basis_change_with_sweep(
-    pauli_strings: list[ops.PauliString], qid_list: list[ops.Qid]
+    pauli_strings: Sequence[ops.PauliString], qid_list: Sequence[ops.Qid]
 ) -> dict[str, float]:
     """Decide single-qubit rotation sweep parameters for basis change.
 
@@ -274,20 +274,28 @@ def _pauli_strings_to_basis_change_with_sweep(
 def _generate_basis_change_circuits(
     normalized_circuits_to_pauli: dict[circuits.FrozenCircuit, list[list[ops.PauliString]]],
     insert_strategy: circuits.InsertStrategy,
+    qubits_to_measure: Sequence[ops.Qid] | None = None,
 ) -> list[circuits.Circuit]:
     """Generates basis change circuits for each group of Pauli strings."""
     pauli_measurement_circuits: list[circuits.Circuit] = []
 
     for input_circuit, pauli_string_groups in normalized_circuits_to_pauli.items():
+        global_qubits = list(qubits_to_measure) if qubits_to_measure is not None else None
+
         basis_change_circuits = []
         input_circuit_unfrozen = input_circuit.unfreeze()
         for pauli_strings in pauli_string_groups:
-            # Extract qubits from Pauli strings
-            qid_list = _extract_readout_qubits(pauli_strings)
+            if global_qubits is not None:
+                # Use the user-provided override
+                current_qid_list = global_qubits
+            else:
+                # Extract qubits from Pauli strings
+                current_qid_list = _extract_readout_qubits(pauli_strings)
+
             basis_change_circuit = circuits.Circuit(
                 input_circuit_unfrozen,
-                _pauli_strings_to_basis_change_ops(pauli_strings, qid_list),
-                ops.measure(*qid_list, key="result"),
+                _pauli_strings_to_basis_change_ops(pauli_strings, current_qid_list),
+                ops.measure(*current_qid_list, key="result"),
                 strategy=insert_strategy,
             )
             basis_change_circuits.append(basis_change_circuit)
@@ -299,28 +307,55 @@ def _generate_basis_change_circuits(
 def _generate_basis_change_circuits_with_sweep(
     normalized_circuits_to_pauli: dict[circuits.FrozenCircuit, list[list[ops.PauliString]]],
     insert_strategy: circuits.InsertStrategy,
+    qubits_to_measure: Sequence[ops.Qid] | None = None,
 ) -> tuple[list[circuits.Circuit], list[cirq.Sweepable]]:
     """Generates basis change circuits for each group of Pauli strings with sweep."""
     parameterized_circuits: list[circuits.Circuit] = []
     sweep_params: list[cirq.Sweepable] = []
     for input_circuit, pauli_string_groups in normalized_circuits_to_pauli.items():
-        for pauli_strings in pauli_string_groups:
-            # Extract qubits from Pauli strings
-            qid_list = _extract_readout_qubits(pauli_strings)
-            phi_symbols = sympy.symbols(f"phi:{len(qid_list)}")
-            theta_symbols = sympy.symbols(f"theta:{len(qid_list)}")
+        # If qubits_to_measure is provided, use it
+        if qubits_to_measure:
+            phi_symbols = sympy.symbols(f"phi:{len(qubits_to_measure)}")
+            theta_symbols = sympy.symbols(f"theta:{len(qubits_to_measure)}")
+
             # Create phased gates and measurement operator
             phased_gates = [
                 ops.PhasedXPowGate(phase_exponent=(a - 1) / 2, exponent=b)(qubit)
-                for a, b, qubit in zip(phi_symbols, theta_symbols, qid_list)
+                for a, b, qubit in zip(phi_symbols, theta_symbols, qubits_to_measure)
             ]
-            measurement_op = ops.M(*qid_list, key="result")
+            measurement_op = ops.M(*qubits_to_measure, key="result")
+
             parameterized_circuit = circuits.Circuit(
                 input_circuit.unfreeze(), phased_gates, measurement_op, strategy=insert_strategy
             )
-            sweep_param = _pauli_strings_to_basis_change_with_sweep(pauli_strings, qid_list)
-            parameterized_circuits.append(parameterized_circuit)
+            sweep_param = []
+            for pauli_strings in pauli_string_groups:
+                sweep_param.append(
+                    _pauli_strings_to_basis_change_with_sweep(pauli_strings, qubits_to_measure)
+                )
             sweep_params.append(sweep_param)
+            parameterized_circuits.append(parameterized_circuit)
+
+        else:
+            for pauli_strings in pauli_string_groups:
+                # Extract qubits from Pauli strings
+                qid_list = _extract_readout_qubits(pauli_strings)
+                phi_symbols = sympy.symbols(f"phi:{len(qid_list)}")
+                theta_symbols = sympy.symbols(f"theta:{len(qid_list)}")
+                # Create phased gates and measurement operator
+                phased_gates = [
+                    ops.PhasedXPowGate(phase_exponent=(a - 1) / 2, exponent=b)(qubit)
+                    for a, b, qubit in zip(phi_symbols, theta_symbols, qid_list)
+                ]
+                measurement_op = ops.M(*qid_list, key="result")
+                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
+                )
+                parameterized_circuits.append(parameterized_circuit)
+                sweep_params.append(sweep_param_dict)
     return parameterized_circuits, sweep_params
 
 
@@ -376,6 +411,7 @@ def _process_pauli_measurement_results(
     pauli_repetitions: int,
     timestamp: float,
     disable_readout_mitigation: bool = False,
+    fixed_calibration_key: tuple[ops.Qid, ...] | None = None,
 ) -> list[PauliStringMeasurementResult]:
     """Calculates both error-mitigated expectation values and unmitigated expectation values
     from measurement results.
@@ -396,22 +432,28 @@ def _process_pauli_measurement_results(
         timestamp: The timestamp of the calibration results.
         disable_readout_mitigation: If set to True, returns no error-mitigated error
             expectation values.
+        fixed_calibration_key: If provided, uses this key to retrieve the calibration result
+            from `calibration_results` for all Pauli strings, regardless of their specific
+            support. This is used when `measure_on_full_support` is True.
 
     Returns:
         A list of PauliStringMeasurementResult.
     """
-
     pauli_measurement_results: list[PauliStringMeasurementResult] = []
 
     for pauli_group_index, circuit_result in enumerate(circuit_results):
         measurement_results = circuit_result.measurements["result"]
         pauli_strs = pauli_string_groups[pauli_group_index]
-        pauli_readout_qubits = _extract_readout_qubits(pauli_strs)
+
+        if fixed_calibration_key is not None:
+            pauli_readout_qubits = list(fixed_calibration_key)
+            calibration_key = fixed_calibration_key
+        else:
+            pauli_readout_qubits = _extract_readout_qubits(pauli_strs)
+            calibration_key = tuple(pauli_readout_qubits)
 
         calibration_result = (
-            calibration_results[tuple(pauli_readout_qubits)]
-            if not disable_readout_mitigation
-            else None
+            calibration_results[calibration_key] if not disable_readout_mitigation else None
         )
 
         for pauli_str in pauli_strs:
@@ -488,6 +530,7 @@ def measure_pauli_strings(
     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.
@@ -521,7 +564,11 @@ def measure_pauli_strings(
         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:
@@ -542,12 +589,24 @@ def measure_pauli_strings(
 
     # Extract unique qubit tuples from input pauli strings
     unique_qubit_tuples = set()
-    for pauli_string_groups in normalized_circuits_to_pauli.values():
-        for pauli_strings in pauli_string_groups:
-            unique_qubit_tuples.add(tuple(_extract_readout_qubits(pauli_strings)))
+    if measure_on_full_support:
+        full_support: set[ops.Qid] = set()
+        for pauli_string_groups in normalized_circuits_to_pauli.values():
+            for pauli_strings in pauli_string_groups:
+                for pauli_string in pauli_strings:
+                    full_support.update(pauli_string.qubits)
+        # One calibration group
+        unique_qubit_tuples.add(tuple(sorted(full_support)))
+    else:
+        for pauli_string_groups in normalized_circuits_to_pauli.values():
+            for pauli_strings in pauli_string_groups:
+                unique_qubit_tuples.add(tuple(_extract_readout_qubits(pauli_strings)))
+
     # qubits_list is a list of qubit tuples
     qubits_list = sorted(unique_qubit_tuples)
 
+    qubits_to_measure_arg = list(qubits_list[0]) if measure_on_full_support else None
+
     # Build the basis-change circuits for each Pauli string group
     pauli_measurement_circuits: list[circuits.Circuit] = []
     sweep_params: list[cirq.Sweepable] = []
@@ -561,7 +620,7 @@ def measure_pauli_strings(
 
     if use_sweep:
         pauli_measurement_circuits, sweep_params = _generate_basis_change_circuits_with_sweep(
-            normalized_circuits_to_pauli, insert_strategy
+            normalized_circuits_to_pauli, insert_strategy, qubits_to_measure_arg
         )
 
         # Run benchmarking using sweep for readout calibration
@@ -578,7 +637,7 @@ def measure_pauli_strings(
 
     else:
         pauli_measurement_circuits = _generate_basis_change_circuits(
-            normalized_circuits_to_pauli, insert_strategy
+            normalized_circuits_to_pauli, insert_strategy, qubits_to_measure_arg
         )
 
         # Run shuffled benchmarking for readout calibration
@@ -595,26 +654,43 @@ def measure_pauli_strings(
     # Process the results to calculate expectation values
     results: list[CircuitToPauliStringsMeasurementResult] = []
     circuit_result_index = 0
+    input_circuit_index = 0
+
     for input_circuit, pauli_string_groups in normalized_circuits_to_pauli.items():
         disable_readout_mitigation = False if num_random_bitstrings != 0 else True
 
         circuits_results_for_group: Sequence[cirq.ResultDict] | Sequence[cirq.Result] = []
-        results_slice = slice(circuit_result_index, circuit_result_index + len(pauli_string_groups))
-        if use_sweep:
-            circuits_results_for_group = [r[0] for r in sweep_circuits_results[results_slice]]
 
+        if use_sweep:
+            if measure_on_full_support:
+                circuits_results_for_group = sweep_circuits_results[input_circuit_index]
+                input_circuit_index += 1
+            else:
+                results_slice = slice(
+                    circuit_result_index, circuit_result_index + len(pauli_string_groups)
+                )
+                circuits_results_for_group = [r[0] for r in sweep_circuits_results[results_slice]]
+                circuit_result_index += len(pauli_string_groups)
         else:
+            results_slice = slice(
+                circuit_result_index, circuit_result_index + len(pauli_string_groups)
+            )
             circuits_results_for_group = circuits_results[results_slice]
+            circuit_result_index += len(pauli_string_groups)
 
-        circuit_result_index += len(pauli_string_groups)
-
+        fixed_calibration_key = (
+            tuple(qubits_to_measure_arg)
+            if measure_on_full_support and qubits_to_measure_arg is not None
+            else None
+        )
         pauli_measurement_results = _process_pauli_measurement_results(
             pauli_string_groups,
             circuits_results_for_group,
             calibration_results,
             pauli_repetitions,
             time.time(),
             disable_readout_mitigation,
+            fixed_calibration_key,
         )
         results.append(
             CircuitToPauliStringsMeasurementResult(

cirq-core/cirq/contrib/paulistring/pauli_string_measurement_with_readout_mitigation_test.py

@@ -494,32 +494,49 @@ def test_many_group_pauli_in_circuits_with_coefficient(use_sweep: bool) -> None:
     circuit_3 = cirq.FrozenCircuit(_create_ghz(8, qubits_3))
 
     circuits_to_pauli: dict[cirq.FrozenCircuit, list[list[cirq.PauliString]]] = {}
+
     circuits_to_pauli[circuit_1] = [
         _generate_qwc_paulis(
-            _generate_random_pauli_string(qubits_1, enable_coeff=True, allow_pauli_i=False), 4
+            _generate_random_pauli_string(qubits_1, enable_coeff=True, allow_pauli_i=False), 2
         )
     ]
+
     circuits_to_pauli[circuit_2] = [
         _generate_qwc_paulis(
-            _generate_random_pauli_string(qubits_2, enable_coeff=True, allow_pauli_i=False), 5
+            _generate_random_pauli_string(qubits_2, enable_coeff=True, allow_pauli_i=False), 2
         )
     ]
+
     circuits_to_pauli[circuit_3] = [
         _generate_qwc_paulis(
-            _generate_random_pauli_string(qubits_3, enable_coeff=True, allow_pauli_i=False), 6
+            _generate_random_pauli_string(qubits_3, enable_coeff=True, allow_pauli_i=False), 2
         )
     ]
 
     sampler = NoisySingleQubitReadoutSampler(p0=0.03, p1=0.05, seed=1234)
     simulator = cirq.Simulator()
 
     circuits_with_pauli_expectations = measure_pauli_strings(
-        circuits_to_pauli, sampler, 300, 300, 300, np.random.default_rng(), use_sweep
+        circuits_to_pauli,
+        sampler,
+        300,
+        300,
+        300,
+        np.random.default_rng(),
+        use_sweep,
+        measure_on_full_support=True,
     )
 
     for circuit_with_pauli_expectations in circuits_with_pauli_expectations:
         assert isinstance(circuit_with_pauli_expectations.circuit, cirq.FrozenCircuit)
 
+        expected_group_count = len(circuits_to_pauli[circuit_with_pauli_expectations.circuit][0])
+
+        assert len(circuit_with_pauli_expectations.results) == expected_group_count, (
+            f"Expected {expected_group_count} results (groups) for circuit, "
+            f"but got {len(circuit_with_pauli_expectations.results)}."
+        )
+
         expected_val_simulation = simulator.simulate(
             circuit_with_pauli_expectations.circuit.unfreeze()
         )