Add parameter sweep support for Pauli string measurements with readout mitigation

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

@@ -18,13 +18,18 @@
 
 import itertools
 import time
-from typing import cast, Sequence, TYPE_CHECKING
+from typing import cast, Sequence, TYPE_CHECKING, Union
 
 import attrs
 import numpy as np
-
-from cirq import circuits, ops, work
-from cirq.contrib.shuffle_circuits import run_shuffled_with_readout_benchmarking
+import sympy
+
+from cirq import circuits, ops, study, work
+from cirq.contrib.shuffle_circuits import (
+    run_shuffled_with_readout_benchmarking,
+    run_sweep_with_readout_benchmarking,
+)
+from cirq.experiments import SingleQubitReadoutCalibrationResult
 from cirq.experiments.readout_confusion_matrix import TensoredConfusionMatrices
 
 if TYPE_CHECKING:
@@ -198,10 +203,10 @@ def _validate_input(
 
 
 def _normalize_input_paulis(
-    circuits_to_pauli: (
-        dict[circuits.FrozenCircuit, list[ops.PauliString]]
-        | dict[circuits.FrozenCircuit, list[list[ops.PauliString]]]
-    ),
+    circuits_to_pauli: Union[
+        dict[circuits.FrozenCircuit, list[ops.PauliString]],
+        dict[circuits.FrozenCircuit, list[list[ops.PauliString]]],
+    ],
 ) -> dict[circuits.FrozenCircuit, list[list[ops.PauliString]]]:
     first_value = next(iter(circuits_to_pauli.values()))
     if (
@@ -233,6 +238,77 @@ def _pauli_strings_to_basis_change_ops(
     return operations
 
 
+def _pauli_strings_to_basis_change_with_sweep(
+    pauli_strings: list[ops.PauliString], qid_list: list[ops.Qid]
+) -> dict[str, float]:
+    """Decide single-qubit rotation sweep parameters for basis change."""
+    params_dict = {}
+
+    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:
+            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
+            elif pauli_op == ops.Y:
+                params_dict[f"phi{qid}"] = 1.0
+                params_dict[f"theta{qid}"] = 1 / 2
+                break
+    return params_dict
+
+
+def _generate_basis_change_circuits(
+    normalized_circuits_to_pauli: dict[circuits.FrozenCircuit, list[list[ops.PauliString]]],
+) -> 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():
+        qid_list = list(sorted(input_circuit.all_qubits()))
+        basis_change_circuits = []
+        input_circuit_unfrozen = input_circuit.unfreeze()
+        for pauli_strings in pauli_string_groups:
+            basis_change_circuit = (
+                input_circuit_unfrozen
+                + _pauli_strings_to_basis_change_ops(pauli_strings, qid_list)
+                + ops.measure(*qid_list, key="m")
+            )
+            basis_change_circuits.append(basis_change_circuit)
+        pauli_measurement_circuits.extend(basis_change_circuits)
+
+    return pauli_measurement_circuits
+
+
+def _generate_basis_change_circuits_with_sweep(
+    normalized_circuits_to_pauli: dict[circuits.FrozenCircuit, list[list[ops.PauliString]]],
+) -> tuple[list[circuits.Circuit], list[study.Sweepable]]:
+    """Generates basis change circuits for each group of Pauli strings with sweep."""
+    parameterized_circuits = list[circuits.Circuit]()
+    sweep_params = list[study.Sweepable]()
+    for input_circuit, pauli_string_groups in normalized_circuits_to_pauli.items():
+        qid_list = list(sorted(input_circuit.all_qubits()))
+        phi_symbols = sympy.symbols(f"phi:{len(qid_list)}")
+        theta_symbols = sympy.symbols(f"theta:{len(qid_list)}")
+
+        parameterized_circuit = input_circuit.unfreeze() + circuits.Circuit(
+            [
+                ops.PhasedXPowGate(phase_exponent=(a - 1) / 2, exponent=b)(qubit)
+                for a, b, qubit in zip(phi_symbols, theta_symbols, qid_list)
+            ],
+            ops.M(*qid_list, key="m"),
+        )
+        sweep_param = []
+        for pauli_strings in pauli_string_groups:
+            sweep_param.append(_pauli_strings_to_basis_change_with_sweep(pauli_strings, qid_list))
+        sweep_params.append(sweep_param)
+        parameterized_circuits.append(parameterized_circuit)
+
+    return parameterized_circuits, sweep_params
+
+
 def _build_one_qubit_confusion_matrix(e0: float, e1: float) -> np.ndarray:
     """Builds a 2x2 confusion matrix for a single qubit.
 
@@ -281,7 +357,7 @@ def _build_many_one_qubits_empty_confusion_matrix(qubits_length: int) -> list[np
 def _process_pauli_measurement_results(
     qubits: list[ops.Qid],
     pauli_string_groups: list[list[ops.PauliString]],
-    circuit_results: list[ResultDict],
+    circuit_results: Sequence[ResultDict],
     calibration_results: dict[tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult],
     pauli_repetitions: int,
     timestamp: float,
@@ -376,7 +452,8 @@ def measure_pauli_strings(
     pauli_repetitions: int,
     readout_repetitions: int,
     num_random_bitstrings: int,
-    rng_or_seed: np.random.Generator | int,
+    rng_or_seed: Union[np.random.Generator, int],
+    use_sweep: bool = False,
 ) -> list[CircuitToPauliStringsMeasurementResult]:
     """Measures expectation values of Pauli strings on given circuits with/without
     readout error mitigation.
@@ -385,7 +462,8 @@ def measure_pauli_strings(
     For each circuit and its associated list of QWC pauli string group, it:
     1.  Constructs circuits to measure the Pauli string expectation value by
         adding basis change moments and measurement operations.
-    2.  Runs shuffled readout benchmarking on these circuits to calibrate readout errors.
+    2.  If `num_random_bitstrings` is greater than zero, performing readout
+        benchmarking (shuffled or sweep-based) to calibrate readout errors.
     3.  Mitigates readout errors using the calibrated confusion matrices.
     4.  Calculates and returns both error-mitigated and unmitigated expectation values for
     each Pauli string.
@@ -406,6 +484,8 @@ def measure_pauli_strings(
         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..
 
     Returns:
         A list of CircuitToPauliStringsMeasurementResult objects, where each object contains:
@@ -435,45 +515,63 @@ def measure_pauli_strings(
 
     # Build the basis-change circuits for each Pauli string group
     pauli_measurement_circuits = list[circuits.Circuit]()
-    for input_circuit, pauli_string_groups in normalized_circuits_to_pauli.items():
-        qid_list = list(sorted(input_circuit.all_qubits()))
-        basis_change_circuits = []
-        input_circuit_unfrozen = input_circuit.unfreeze()
-        for pauli_strings in pauli_string_groups:
-            basis_change_circuit = (
-                input_circuit_unfrozen
-                + _pauli_strings_to_basis_change_ops(pauli_strings, qid_list)
-                + ops.measure(*qid_list, key="m")
-            )
-            basis_change_circuits.append(basis_change_circuit)
-        pauli_measurement_circuits.extend(basis_change_circuits)
+    sweep_params = list[study.Sweepable]()
+    circuits_results: Union[Sequence[ResultDict], Sequence[Sequence[study.Result]]] = []
+    # calibration_results = list[Dict[Tuple[ops.Qid, ...], SingleQubitReadoutCalibrationResult]]()
 
-    # Run shuffled benchmarking for readout calibration
-    circuits_results, calibration_results = run_shuffled_with_readout_benchmarking(
-        input_circuits=pauli_measurement_circuits,
-        sampler=sampler,
-        circuit_repetitions=pauli_repetitions,
-        rng_or_seed=rng_or_seed,
-        qubits=[list(qubits) for qubits in qubits_list],
-        num_random_bitstrings=num_random_bitstrings,
-        readout_repetitions=readout_repetitions,
-    )
+    if use_sweep:
+        pauli_measurement_circuits, sweep_params = _generate_basis_change_circuits_with_sweep(
+            normalized_circuits_to_pauli
+        )
+
+        # Run benchmarking using sweep for readout calibration
+        circuits_results, calibration_results = run_sweep_with_readout_benchmarking(
+            input_circuits=pauli_measurement_circuits,
+            sweep_params=sweep_params,
+            sampler=sampler,
+            circuit_repetitions=pauli_repetitions,
+            rng_or_seed=rng_or_seed,
+            qubits=[list(qubits) for qubits in qubits_list],
+            num_random_bitstrings=num_random_bitstrings,
+            readout_repetitions=readout_repetitions,
+        )
+
+    else:
+        pauli_measurement_circuits = _generate_basis_change_circuits(normalized_circuits_to_pauli)
+
+        # Run shuffled benchmarking for readout calibration
+        circuits_results, calibration_results = run_shuffled_with_readout_benchmarking(
+            input_circuits=pauli_measurement_circuits,
+            sampler=sampler,
+            circuit_repetitions=pauli_repetitions,
+            rng_or_seed=rng_or_seed,
+            qubits=[list(qubits) for qubits in qubits_list],
+            num_random_bitstrings=num_random_bitstrings,
+            readout_repetitions=readout_repetitions,
+        )
 
     # Process the results to calculate expectation values
     results: list[CircuitToPauliStringsMeasurementResult] = []
     circuit_result_index = 0
-    for input_circuit, pauli_string_groups in normalized_circuits_to_pauli.items():
+    for i, (input_circuit, pauli_string_groups) in enumerate(normalized_circuits_to_pauli.items()):
 
         qubits_in_circuit = tuple(sorted(input_circuit.all_qubits()))
 
         disable_readout_mitigation = False if num_random_bitstrings != 0 else True
 
+        circuits_results_for_group: Union[ResultDict, Sequence[study.Result]] = []
+        if use_sweep:
+            circuits_results_for_group = circuits_results[i]
+        else:
+            circuits_results_for_group = circuits_results[
+                circuit_result_index : circuit_result_index + len(pauli_string_groups)
+            ]
+            circuit_result_index += len(pauli_string_groups)
+
         pauli_measurement_results = _process_pauli_measurement_results(
             list(qubits_in_circuit),
             pauli_string_groups,
-            circuits_results[
-                circuit_result_index : circuit_result_index + len(pauli_string_groups)
-            ],
+            circuits_results_for_group,
             calibration_results,
             pauli_repetitions,
             time.time(),
@@ -485,5 +583,4 @@ def measure_pauli_strings(
             )
         )
 
-        circuit_result_index += len(pauli_string_groups)
     return results