Update RB to reduce the size of the generated circuits

cirq-core/cirq/experiments/qubit_characterizations.py

@@ -17,8 +17,9 @@
 import dataclasses
 import functools
 import itertools
-from typing import Any, cast, Iterator, Mapping, Sequence, TYPE_CHECKING
+from typing import Any, cast, Iterator, Mapping, overload, Sequence, TYPE_CHECKING
 
+import attrs
 import numpy as np
 from matplotlib import pyplot as plt
 
@@ -28,6 +29,7 @@
 import cirq.vis.heatmap as cirq_heatmap
 import cirq.vis.histogram as cirq_histogram
 from cirq import circuits, ops, protocols
+from cirq._compat import deprecated
 from cirq.devices import grid_qubit
 
 if TYPE_CHECKING:
@@ -36,6 +38,12 @@
     import cirq
 
 
+def _canonize_clifford_sequences(
+    sequences: list[list[ops.SingleQubitCliffordGate]],
+) -> list[list[ops.SingleQubitCliffordGate]]:
+    return [[_reduce_gate_seq(seq)] for seq in sequences]
+
+
 @dataclasses.dataclass
 class Cliffords:
     """The single-qubit Clifford group, decomposed into elementary gates.
@@ -333,6 +341,38 @@ def plot(
         return axes
 
 
+@attrs.frozen
+class RBParameters:
+    r"""Parameters for running randomized benchmarking.
+
+    Arguments:
+        num_clifford_range: The different numbers of Cliffords in the RB study.
+        num_circuits: The number of random circuits generated for each
+            number of Cliffords.
+        repetitions: The number of repetitions of each circuit.
+        use_xy_basis: Determines if the Clifford gates are built with x and y
+            rotations (True) or x and z rotations (False).
+        strict_basis: whether to use only Cliffords that can be represented by at
+            most 2 gates of the chosen basis. For example,
+            if True and use_xy_basis is True, this excludes $I, Z, \sqrt(Z), \-sqrt(Z)^\dagger$.
+            if True and use_xy_basis is False, this excludes $I, Y, \sqrt(Y), -\sqrt(Y)^\dagger$.
+    """
+
+    num_clifford_range: Sequence[int] = tuple(np.logspace(np.log10(5), 3, 5, dtype=int))
+    num_circuits: int = 10
+    repetitions: int = 600
+    use_xy_basis: bool = False
+    strict_basis: bool = True
+
+    def gateset(self) -> list[list[ops.SingleQubitCliffordGate]]:
+        clifford_group = _single_qubit_cliffords()
+        sequences = clifford_group.c1_in_xy if self.use_xy_basis else clifford_group.c1_in_xz
+        if self.strict_basis:
+            sequences = [seq for seq in sequences if len(seq) == 2]
+        return _canonize_clifford_sequences(sequences)
+
+
+@deprecated(deadline='v2.0', fix='please use single_qubit_rb instead')
 def single_qubit_randomized_benchmarking(
     sampler: cirq.Sampler,
     qubit: cirq.Qid,
@@ -376,17 +416,20 @@ def single_qubit_randomized_benchmarking(
         A RandomizedBenchMarkResult object that stores and plots the result.
     """
 
-    result = parallel_single_qubit_randomized_benchmarking(
+    return parallel_single_qubit_rb(
         sampler,
-        (qubit,),
-        use_xy_basis,
-        num_clifford_range=num_clifford_range,
-        num_circuits=num_circuits,
-        repetitions=repetitions,
+        qubit,
+        RBParameters(
+            num_clifford_range=num_clifford_range,
+            num_circuits=num_circuits,
+            repetitions=repetitions,
+            use_xy_basis=use_xy_basis,
+            strict_basis=False,
+        ),
     )
-    return result.results_dictionary[qubit]
 
 
+@deprecated(deadline='v2.0', fix='please use single_qubit_rb instead')
 def parallel_single_qubit_randomized_benchmarking(
     sampler: cirq.Sampler,
     qubits: Sequence[cirq.Qid],
@@ -413,35 +456,93 @@ def parallel_single_qubit_randomized_benchmarking(
         num_circuits: The number of random circuits generated for each
             number of Cliffords.
         repetitions: The number of repetitions of each circuit.
+    Returns:
+        A dictionary from qubits to RandomizedBenchMarkResult objects.
+    """
+    return parallel_single_qubit_rb(
+        sampler,
+        qubits,
+        RBParameters(
+            num_clifford_range=num_clifford_range,
+            num_circuits=num_circuits,
+            repetitions=repetitions,
+            use_xy_basis=use_xy_basis,
+            strict_basis=False,
+        ),
+    )
+
 
+@overload
+def parallel_single_qubit_rb(
+    sampler: cirq.Sampler,
+    qubits: cirq.Qid,
+    parameters: RBParameters = RBParameters(),
+    rng_or_seed: np.random.Generator | int | None = None,
+) -> RandomizedBenchMarkResult: ...
+
+
+@overload
+def parallel_single_qubit_rb(
+    sampler: cirq.Sampler,
+    qubits: Sequence[cirq.Qid],
+    parameters: RBParameters = RBParameters(),
+    rng_or_seed: np.random.Generator | int | None = None,
+) -> ParallelRandomizedBenchmarkingResult: ...
+
+
+def parallel_single_qubit_rb(
+    sampler: cirq.Sampler,
+    qubits: Sequence[cirq.Qid] | cirq.Qid,
+    parameters: RBParameters = RBParameters(),
+    rng_or_seed: np.random.Generator | int | None = None,
+) -> RandomizedBenchMarkResult | ParallelRandomizedBenchmarkingResult:
+    """Clifford-based randomized benchmarking (RB) single qubits in parallel.
+
+    Args:
+        sampler: The quantum engine or simulator to run the circuits.
+        qubits: The qubit(s) to benchmark.
+        parameters: The parameters of the experiment.
+        rng_or_seed: A np.random.Generator object or seed.
     Returns:
         A dictionary from qubits to RandomizedBenchMarkResult objects.
     """
+    is_single_qubit = False
+    if isinstance(qubits, ops.Qid):
+        qubits = (qubits,)
+        is_single_qubit = True
+
+    rng_or_seed = (
+        rng_or_seed
+        if isinstance(rng_or_seed, np.random.Generator)
+        else np.random.default_rng(rng_or_seed)
+    )
 
-    clifford_group = _single_qubit_cliffords()
-    c1 = clifford_group.c1_in_xy if use_xy_basis else clifford_group.c1_in_xz
+    c1 = parameters.gateset()
 
     # create circuits
     circuits_all: list[cirq.AbstractCircuit] = []
-    for num_cliffords in num_clifford_range:
-        for _ in range(num_circuits):
-            circuits_all.append(_create_parallel_rb_circuit(qubits, num_cliffords, c1))
+    for num_cliffords in parameters.num_clifford_range:
+        for _ in range(parameters.num_circuits):
+            circuits_all.append(_create_parallel_rb_circuit(qubits, num_cliffords, c1, rng_or_seed))
 
     # run circuits
-    results = sampler.run_batch(circuits_all, repetitions=repetitions)
+    results = sampler.run_batch(circuits_all, repetitions=parameters.repetitions)
     gnd_probs: dict = {q: [] for q in qubits}
     idx = 0
-    for num_cliffords in num_clifford_range:
+    for num_cliffords in parameters.num_clifford_range:
         excited_probs: dict[cirq.Qid, list[float]] = {q: [] for q in qubits}
-        for _ in range(num_circuits):
+        for _ in range(parameters.num_circuits):
             result = results[idx][0]
             for qubit in qubits:
                 excited_probs[qubit].append(np.mean(result.measurements[str(qubit)]))
             idx += 1
         for qubit in qubits:
             gnd_probs[qubit].append(1.0 - np.mean(excited_probs[qubit]))
+
+    if is_single_qubit:
+        return RandomizedBenchMarkResult(parameters.num_clifford_range, gnd_probs[qubits[0]])
     return ParallelRandomizedBenchmarkingResult(
-        {q: RandomizedBenchMarkResult(num_clifford_range, gnd_probs[q]) for q in qubits}
+        {q: RandomizedBenchMarkResult(parameters.num_clifford_range, gnd_probs[q]) for q in qubits}
     )
 
 
@@ -677,9 +778,14 @@ def _measurement(two_qubit_circuit: circuits.Circuit) -> np.ndarray:
 
 
 def _create_parallel_rb_circuit(
-    qubits: Sequence[cirq.Qid], num_cliffords: int, c1: list
+    qubits: Sequence[cirq.Qid],
+    num_cliffords: int,
+    c1: list[list[ops.SingleQubitCliffordGate]],
+    rng: np.random.Generator | None = None,
 ) -> cirq.Circuit:
-    sequences_to_zip = [_random_single_q_clifford(qubit, num_cliffords, c1) for qubit in qubits]
+    sequences_to_zip = [
+        _random_single_q_clifford(qubit, num_cliffords, c1, rng) for qubit in qubits
+    ]
     # Ensure each sequence has the same number of moments.
     num_moments = max(len(sequence) for sequence in sequences_to_zip)
     for q, sequence in zip(qubits, sequences_to_zip):
@@ -730,11 +836,14 @@ def _two_qubit_clifford_matrices(q_0: cirq.Qid, q_1: cirq.Qid, cliffords: Cliffo
 
 
 def _random_single_q_clifford(
-    qubit: cirq.Qid, num_cfds: int, cfds: Sequence[Sequence[cirq.ops.SingleQubitCliffordGate]]
+    qubit: cirq.Qid,
+    num_cfds: int,
+    cfds: Sequence[Sequence[cirq.ops.SingleQubitCliffordGate]],
+    rng: np.random.Generator | None = None,
 ) -> list[cirq.Operation]:
-    clifford_group_size = 24
     operations = [[gate.to_phased_xz_gate()(qubit) for gate in gates] for gates in cfds]
-    gate_ids = np.random.choice(clifford_group_size, num_cfds).tolist()
+    choice_fn = rng.choice if rng else np.random.choice
+    gate_ids = choice_fn(len(cfds), num_cfds).tolist()
     adjoint = _reduce_gate_seq([gate for gate_id in gate_ids for gate in cfds[gate_id]]) ** -1
     return [op for gate_id in gate_ids for op in operations[gate_id]] + [
         adjoint.to_phased_xz_gate()(qubit)

cirq-core/cirq/experiments/qubit_characterizations_test.py

@@ -14,6 +14,9 @@
 
 from __future__ import annotations
 
+import os
+from unittest import mock
+
 import matplotlib.pyplot as plt
 import numpy as np
 import pytest
@@ -104,6 +107,7 @@ def check_distinct(unitaries):
         assert num_x <= 1
 
 
+@mock.patch.dict(os.environ, clear='CIRQ_TESTING')
 def test_single_qubit_randomized_benchmarking():
     # Check that the ground state population at the end of the Clifford
     # sequences is always unity.
@@ -116,7 +120,8 @@ def test_single_qubit_randomized_benchmarking():
     assert np.isclose(results.pauli_error(), 0.0, atol=1e-7)  # warning is expected
 
 
-def test_parallel_single_qubit_randomized_benchmarking():
+@mock.patch.dict(os.environ, clear='CIRQ_TESTING')
+def test_parallel_single_qubit_parallel_single_qubit_randomized_benchmarking():
     # Check that the ground state population at the end of the Clifford
     # sequences is always unity.
     simulator = sim.Simulator()
@@ -229,13 +234,20 @@ def test_tomography_plot_raises_for_incorrect_number_of_axes():
         result.plot(axes)
 
 
-def test_single_qubit_cliffords_gateset():
+@pytest.mark.parametrize('num_cliffords', range(5, 10))
+@pytest.mark.parametrize('use_xy_basis', [False, True])
+@pytest.mark.parametrize('strict_basis', [False, True])
+def test_single_qubit_cliffords_gateset(num_cliffords, use_xy_basis, strict_basis):
     qubits = [GridQubit(0, i) for i in range(4)]
-    clifford_group = cirq.experiments.qubit_characterizations._single_qubit_cliffords()
+    c1_in_xy = cirq.experiments.qubit_characterizations.RBParameters(
+        use_xy_basis=use_xy_basis, strict_basis=strict_basis
+    ).gateset()
     c = cirq.experiments.qubit_characterizations._create_parallel_rb_circuit(
-        qubits, 5, clifford_group.c1_in_xy
+        qubits, num_cliffords, c1_in_xy
     )
     device = cirq.testing.ValidatingTestDevice(
         qubits=qubits, allowed_gates=(cirq.ops.PhasedXZGate, cirq.MeasurementGate)
     )
     device.validate_circuit(c)
+
+    assert len(c) == num_cliffords + 2