Add ability to pass tags to XEB and related methods (quantumlib#6836)

Author: NoureldinYosri

cirq-core/cirq/experiments/random_quantum_circuit_generation.py

@@ -250,6 +250,7 @@ def generate_library_of_2q_circuits(
     q0: 'cirq.Qid' = devices.LineQubit(0),
     q1: 'cirq.Qid' = devices.LineQubit(1),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    tags: Sequence[Any] = (),
 ) -> List['cirq.Circuit']:
     """Generate a library of two-qubit Circuits.
 
@@ -266,6 +267,7 @@ def generate_library_of_2q_circuits(
         q0: The first qubit to use when constructing the circuits.
         q1: The second qubit to use when constructing the circuits
         random_state: A random state or seed used to deterministically sample the random circuits.
+        tags: Tags to add to the two qubit operations.
     """
     rs = value.parse_random_state(random_state)
     exponents = np.linspace(0, 7 / 4, 8)
@@ -278,7 +280,7 @@ def generate_library_of_2q_circuits(
             q0,
             q1,
             depth=max_cycle_depth,
-            two_qubit_op_factory=lambda a, b, _: two_qubit_gate(a, b),
+            two_qubit_op_factory=lambda a, b, _: two_qubit_gate(a, b).with_tags(*tags),
             single_qubit_gates=single_qubit_gates,
             seed=rs,
         )

cirq-core/cirq/experiments/random_quantum_circuit_generation_test.py

@@ -457,3 +457,20 @@ def add_pair(neighbor: 'cirq.GridQubit'):
         add_pair(cirq.GridQubit(qubit.row + 1, qubit.col))
 
     return pairs
+
+
+def test_generate_library_of_2q_circuits_with_tags():
+    circuits = generate_library_of_2q_circuits(
+        n_library_circuits=5,
+        two_qubit_gate=cirq.FSimGate(3, 4),
+        max_cycle_depth=13,
+        random_state=9,
+        tags=('test_tag',),
+    )
+    assert len(circuits) == 5
+    for circuit in circuits:
+        for op in circuit.all_operations():
+            if cirq.num_qubits(op) == 1:
+                continue
+            assert op.tags == ('test_tag',)
+            assert op.gate == cirq.FSimGate(3, 4)

cirq-core/cirq/experiments/two_qubit_xeb.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """Provides functions for running and analyzing two-qubit XEB experiments."""
-from typing import Sequence, TYPE_CHECKING, Optional, Tuple, Dict, cast, Mapping
+from typing import Sequence, TYPE_CHECKING, Optional, Tuple, Dict, cast, Mapping, Any
 
 from dataclasses import dataclass
 from types import MappingProxyType
@@ -402,6 +402,7 @@ def parallel_xeb_workflow(
     pairs: Optional[Sequence[tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
     pool: Optional['multiprocessing.pool.Pool'] = None,
     batch_size: int = 9,
+    tags: Sequence[Any] = (),
     **plot_kwargs,
 ) -> Tuple[pd.DataFrame, Sequence['cirq.Circuit'], pd.DataFrame]:
     """A utility method that runs the full XEB workflow.
@@ -422,6 +423,7 @@ def parallel_xeb_workflow(
         batch_size: We call `run_batch` on the sampler, which can speed up execution in certain
             environments. The number of (circuit, cycle_depth) tasks to be run in each batch
             is given by this number.
+        tags: Tags to add to two qubit operations.
         **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
 
     Returns:
@@ -450,6 +452,7 @@ def parallel_xeb_workflow(
         two_qubit_gate=entangling_gate,
         random_state=rs,
         max_cycle_depth=max(cycle_depths),
+        tags=tags,
     )
 
     combs_by_layer = rqcg.get_random_combinations_for_device(
@@ -488,6 +491,7 @@ def parallel_two_qubit_xeb(
     ax: Optional[plt.Axes] = None,
     pairs: Optional[Sequence[tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
     batch_size: int = 9,
+    tags: Sequence[Any] = (),
     **plot_kwargs,
 ) -> TwoQubitXEBResult:
     """A convenience method that runs the full XEB workflow.
@@ -507,6 +511,7 @@ def parallel_two_qubit_xeb(
         batch_size: We call `run_batch` on the sampler, which can speed up execution in certain
             environments. The number of (circuit, cycle_depth) tasks to be run in each batch
             is given by this number.
+        tags: Tags to add to two qubit operations.
         **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
     Returns:
         A TwoQubitXEBResult object representing the results of the experiment.
@@ -525,6 +530,7 @@ def parallel_two_qubit_xeb(
         random_state=random_state,
         ax=ax,
         batch_size=batch_size,
+        tags=tags,
         **plot_kwargs,
     )
     return TwoQubitXEBResult(fit_exponential_decays(fids))
@@ -544,6 +550,7 @@ def run_rb_and_xeb(
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     pairs: Optional[Sequence[tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
     batch_size: int = 9,
+    tags: Sequence[Any] = (),
 ) -> InferredXEBResult:
     """A convenience method that runs both RB and XEB workflows.
 
@@ -561,6 +568,7 @@ def run_rb_and_xeb(
         batch_size: We call `run_batch` on the sampler, which can speed up execution in certain
             environments. The number of (circuit, cycle_depth) tasks to be run in each batch
             is given by this number.
+        tags: Tags to add to two qubit operations.
 
     Returns:
         An InferredXEBResult object representing the results of the experiment.
@@ -590,6 +598,7 @@ def run_rb_and_xeb(
         n_combinations=xeb_combinations,
         random_state=random_state,
         batch_size=batch_size,
+        tags=tags,
     )
 
     return InferredXEBResult(rb, xeb)

cirq-core/cirq/experiments/z_phase_calibration.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """Provides a method to do z-phase calibration for excitation-preserving gates."""
-from typing import Union, Optional, Sequence, Tuple, Dict, TYPE_CHECKING
+from typing import Union, Optional, Sequence, Tuple, Dict, TYPE_CHECKING, Any
 import multiprocessing
 import multiprocessing.pool
 
@@ -41,6 +41,8 @@ def z_phase_calibration_workflow(
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
     num_workers_or_pool: Union[int, 'multiprocessing.pool.Pool'] = -1,
+    pairs: Optional[Sequence[Tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
+    tags: Sequence[Any] = (),
 ) -> Tuple[xeb_fitting.XEBCharacterizationResult, 'pd.DataFrame']:
     """Perform z-phase calibration for excitation-preserving gates.
 
@@ -77,6 +79,8 @@ def z_phase_calibration_workflow(
             A zero value means no multiprocessing.
             A positive integer value will create a pool with the given number of workers.
             A negative value will create pool with maximum number of workers.
+        pairs: Pairs to use. If not specified, use all pairs between adjacent qubits.
+        tags: Tags to add to two qubit operations.
     Returns:
         - An `XEBCharacterizationResult` object that contains the calibration result.
         - A `pd.DataFrame` comparing the before and after fidelities.
@@ -100,6 +104,8 @@ def z_phase_calibration_workflow(
         n_combinations=n_combinations,
         random_state=random_state,
         pool=pool,
+        tags=tags,
+        pairs=pairs,
     )
 
     if options is None:
@@ -148,6 +154,8 @@ def calibrate_z_phases(
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
     num_workers_or_pool: Union[int, 'multiprocessing.pool.Pool'] = -1,
+    pairs: Optional[Sequence[Tuple['cirq.GridQubit', 'cirq.GridQubit']]] = None,
+    tags: Sequence[Any] = (),
 ) -> Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate']:
     """Perform z-phase calibration for excitation-preserving gates.
 
@@ -184,6 +192,8 @@ def calibrate_z_phases(
             A zero value means no multiprocessing.
             A positive integer value will create a pool with the given number of workers.
             A negative value will create pool with maximum number of workers.
+        pairs: Pairs to use. If not specified, use all pairs between adjacent qubits.
+        tags: Tags to add to two qubit operations.
 
     Returns:
         - A dictionary mapping qubit pairs to the calibrated PhasedFSimGates.
@@ -210,6 +220,8 @@ def calibrate_z_phases(
         random_state=random_state,
         atol=atol,
         num_workers_or_pool=num_workers_or_pool,
+        tags=tags,
+        pairs=pairs,
     )
 
     gates = {}