[XEB] Cycle depths during analysis (#4278)

 - Make the argument optional in benchmark_
 - If provided, you can do a subset of those in sampled_df
 - More robust checking that it matches sampled_df

Author: mpharrigan

cirq-core/cirq/experiments/xeb_fitting.py

@@ -52,7 +52,7 @@
 def benchmark_2q_xeb_fidelities(
     sampled_df: pd.DataFrame,
     circuits: Sequence['cirq.Circuit'],
-    cycle_depths: Sequence[int],
+    cycle_depths: Optional[Sequence[int]] = None,
     param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
     pool: Optional['multiprocessing.pool.Pool'] = None,
 ) -> pd.DataFrame:
@@ -66,18 +66,37 @@ def benchmark_2q_xeb_fidelities(
          sampled_df: The sampled results to benchmark. This is likely produced by a call to
             `sample_2q_xeb_circuits`.
         circuits: The library of circuits corresponding to the sampled results in `sampled_df`.
-        cycle_depths: The sequence of cycle depths to simulate the circuits.
+        cycle_depths: The sequence of cycle depths to benchmark the circuits. If not provided,
+            we use the cycle depths found in `sampled_df`. All requested `cycle_depths` must be
+            present in `sampled_df`.
         param_resolver: If circuits contain parameters, resolve according to this ParamResolver
             prior to simulation
         pool: If provided, execute the simulations in parallel.
 
     Returns:
         A DataFrame with columns 'cycle_depth' and 'fidelity'.
     """
+    sampled_cycle_depths = (
+        sampled_df.index.get_level_values('cycle_depth').drop_duplicates().sort_values()
+    )
+    if cycle_depths is not None:
+        if len(cycle_depths) == 0:
+            raise ValueError("`cycle_depths` should be a non-empty array_like")
+        not_in_sampled = np.setdiff1d(cycle_depths, sampled_cycle_depths)
+        if len(not_in_sampled) > 0:
+            raise ValueError(
+                f"The `cycle_depths` provided include some not "
+                f"available in `sampled_df`: {not_in_sampled}"
+            )
+        sim_cycle_depths = cycle_depths
+    else:
+        sim_cycle_depths = sampled_cycle_depths
     simulated_df = simulate_2q_xeb_circuits(
-        circuits=circuits, cycle_depths=cycle_depths, param_resolver=param_resolver, pool=pool
+        circuits=circuits, cycle_depths=sim_cycle_depths, param_resolver=param_resolver, pool=pool
     )
-    df = sampled_df.join(simulated_df)
+    # Join the `pure_probs` onto `sampled_df`. By using 'inner', we let
+    # the `cycle_depths` argument to this function control what cycle depths are benchmarked.
+    df = sampled_df.join(simulated_df, how='inner').reset_index()
 
     D = 4  # two qubits
     pure_probs = np.array(df['pure_probs'].to_list())
@@ -117,7 +136,7 @@ def _try_keep(k):
     else:
         groupby_names = ['cycle_depth']
 
-    return df.reset_index().groupby(groupby_names).apply(per_cycle_depth).reset_index()
+    return df.groupby(groupby_names).apply(per_cycle_depth).reset_index()
 
 
 class XEBCharacterizationOptions(ABC):

cirq-core/cirq/experiments/xeb_fitting_test.py

@@ -37,31 +37,61 @@
 from cirq.experiments.xeb_sampling import sample_2q_xeb_circuits
 
 
-def test_benchmark_2q_xeb_fidelities():
+@pytest.fixture(scope='module')
+def circuits_cycle_depths_sampled_df():
     q0, q1 = cirq.LineQubit.range(2)
     circuits = [
         rqcg.random_rotations_between_two_qubit_circuit(
             q0, q1, depth=50, two_qubit_op_factory=lambda a, b, _: SQRT_ISWAP(a, b), seed=52
         )
         for _ in range(2)
     ]
-    cycle_depths = np.arange(3, 50, 9)
+    cycle_depths = np.arange(10, 40 + 1, 10)
 
     sampled_df = sample_2q_xeb_circuits(
         sampler=cirq.Simulator(seed=53), circuits=circuits, cycle_depths=cycle_depths
     )
-    fid_df = benchmark_2q_xeb_fidelities(sampled_df, circuits, cycle_depths)
+    return circuits, cycle_depths, sampled_df
+
+
+@pytest.mark.parametrize('pass_cycle_depths', (True, False))
+def test_benchmark_2q_xeb_fidelities(circuits_cycle_depths_sampled_df, pass_cycle_depths):
+    circuits, cycle_depths, sampled_df = circuits_cycle_depths_sampled_df
+
+    if pass_cycle_depths:
+        fid_df = benchmark_2q_xeb_fidelities(sampled_df, circuits, cycle_depths)
+    else:
+        fid_df = benchmark_2q_xeb_fidelities(sampled_df, circuits)
     assert len(fid_df) == len(cycle_depths)
-    for _, row in fid_df.iterrows():
-        assert row['cycle_depth'] in cycle_depths
-        assert row['fidelity'] > 0.98
+    assert sorted(fid_df['cycle_depth'].unique()) == cycle_depths.tolist()
+    assert np.all(fid_df['fidelity'] > 0.98)
 
     fit_df = fit_exponential_decays(fid_df)
     for _, row in fit_df.iterrows():
         assert list(row['cycle_depths']) == list(cycle_depths)
         assert len(row['fidelities']) == len(cycle_depths)
 
 
+def test_benchmark_2q_xeb_subsample_depths(circuits_cycle_depths_sampled_df):
+    circuits, _, sampled_df = circuits_cycle_depths_sampled_df
+    cycle_depths = [10, 20]
+    fid_df = benchmark_2q_xeb_fidelities(sampled_df, circuits, cycle_depths)
+    assert len(fid_df) == len(cycle_depths)
+    assert sorted(fid_df['cycle_depth'].unique()) == cycle_depths
+
+    cycle_depths = [11, 21]
+    with pytest.raises(ValueError):
+        _ = benchmark_2q_xeb_fidelities(sampled_df, circuits, cycle_depths)
+
+    cycle_depths = [10, 100_000]
+    with pytest.raises(ValueError):
+        _ = benchmark_2q_xeb_fidelities(sampled_df, circuits, cycle_depths)
+
+    cycle_depths = []
+    with pytest.raises(ValueError):
+        _ = benchmark_2q_xeb_fidelities(sampled_df, circuits, cycle_depths)
+
+
 def _gridqubits_to_graph_device(qubits: Iterable[cirq.GridQubit]):
     # cirq contrib: routing.gridqubits_to_graph_device
     def _manhattan_distance(qubit1: cirq.GridQubit, qubit2: cirq.GridQubit) -> int: