From ade51ee706c4ba3d45750775ffcf3002c7ed48fc Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 2 Aug 2024 11:01:46 -0700
Subject: [PATCH 01/22] checkpoint

---
 cirq-core/cirq/experiments/xeb_fitting.py     |   8 +-
 .../cirq/experiments/z_phase_calibration.py   | 193 ++++++++++++++++++
 2 files changed, 198 insertions(+), 3 deletions(-)
 create mode 100644 cirq-core/cirq/experiments/z_phase_calibration.py

diff --git a/cirq-core/cirq/experiments/xeb_fitting.py b/cirq-core/cirq/experiments/xeb_fitting.py
index 7f46d2d7f92..55935450d67 100644
--- a/cirq-core/cirq/experiments/xeb_fitting.py
+++ b/cirq-core/cirq/experiments/xeb_fitting.py
@@ -14,7 +14,7 @@
 """Estimation of fidelity associated with experimental circuit executions."""
 import dataclasses
 from abc import abstractmethod, ABC
-from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 import pandas as pd
@@ -385,7 +385,9 @@ def SqrtISwapXEBOptions(*args, **kwargs):
 
 
 def parameterize_circuit(
-    circuit: 'cirq.Circuit', options: XEBCharacterizationOptions
+    circuit: 'cirq.Circuit', options: XEBCharacterizationOptions,     target: Union[ops.GateFamily, ops.Gateset] = ops.Gateset(
+        ops.PhasedFSimGate, ops.ISwapPowGate, ops.FSimGate
+    ),
 ) -> 'cirq.Circuit':
     """Parameterize PhasedFSim-like gates in a given circuit according to
     `phased_fsim_options`.
@@ -393,7 +395,7 @@ def parameterize_circuit(
     gate = options.get_parameterized_gate()
     return circuits.Circuit(
         circuits.Moment(
-            gate.on(*op.qubits) if options.should_parameterize(op) else op
+            gate.on(*op.qubits) if op in target else op
             for op in moment.operations
         )
         for moment in circuit.moments
diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
new file mode 100644
index 00000000000..bd4052d5828
--- /dev/null
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -0,0 +1,193 @@
+# Copyright 2024 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Provides a method to do z-phase calibration for excitation-preserving gates."""
+from typing import Optional, Sequence, Union, Tuple, Dict, TYPE_CHECKING
+import multiprocessing
+import concurrent.futures
+
+import numpy as np
+
+from cirq.experiments import xeb_fitting
+from cirq.experiments.two_qubit_xeb import parallel_xeb_workflow
+from cirq import ops
+
+if TYPE_CHECKING:
+    import cirq
+    import pandas as pd
+
+
+def z_phase_calibration_workflow(
+    sampler: 'cirq.Sampler',
+    qubits: Optional['cirq.GridQubit'] = None,
+    two_qubit_gate: 'cirq.Gate' = ops.CZ,
+    options: Optional[xeb_fitting.XEBPhasedFSimCharacterizationOptions] = None,
+    n_repetitions: int = 10**4,
+    n_combinations: int = 10,
+    n_circuits: int = 20,
+    cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
+    random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    atol: float = 1e-3,
+    pool: Optional[Union[multiprocessing.Pool, concurrent.futures.ThreadPoolExecutor]] = None,
+) -> Tuple[xeb_fitting.XEBCharacterizationResult, 'pd.DataFrame']:
+    """Perform z-phase calibration for excitation-preserving gates.
+
+    For a given excitation-preserving two-qubit gate we assume an error model that can be described
+    using Z-rotations:
+                0: ───Rz(a)───two_qubit_gate───Rz(c)───
+                                │
+                1: ───Rz(b)───two_qubit_gate───Rz(d)───
+    for some angles a, b, c, and d.
+
+    Since the two-qubit gate is a excitation-preserving-gate, it can be represented by an FSimGate
+    and the effect of rotations turns it into a PhasedFSimGate. Using XEB-data we find the
+    PhasedFSimGate parameters that minimize the infidelity of the gate.
+
+    References:
+        - https://arxiv.org/abs/2001.08343
+        - https://arxiv.org/abs/2010.07965
+        - https://arxiv.org/abs/1910.11333
+
+    Args:
+        sampler: The quantum engine or simulator to run the circuits.
+        qubits: Qubits to use. If none, use all qubits on the sampler's device.
+        two_qubit_gate: The entangling gate to use.
+        options: The XEB-fitting options. If None, calibrate all 5 PhasedFSimGate parameters,
+            using the representation of a two-qubit gate as an FSimGate for the initial guess.
+        n_repetitions: The number of repetitions to use.
+        n_combinations: The number of combinations to generate.
+        n_circuits: The number of circuits to generate.
+        cycle_depths: The cycle depths to use.
+        random_state: The random state to use.
+        atol: Absolute tolerance to be used by the minimizer.
+        pool: Optional multi-threading or multi-processing pool.
+
+    Returns:
+        - An `XEBCharacterizationResult` object that contains the calibration result.
+        - A `pd.DataFrame` comparing the before and after fidilities.
+    """
+
+    fids_df_0, circuits, sampled_df = parallel_xeb_workflow(
+        sampler=sampler,
+        qubits=qubits,
+        entangling_gate=two_qubit_gate,
+        n_repetitions=n_repetitions,
+        cycle_depths=cycle_depths,
+        n_circuits=n_circuits,
+        n_combinations=n_combinations,
+        random_state=random_state,
+    )
+
+    if options is None:
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(characterize_theta=False, characterize_phi=False).with_defaults_from_gate(
+            two_qubit_gate
+        )
+
+    p_circuits = [
+        xeb_fitting.parameterize_circuit(circuit, options, ops.GateFamily(two_qubit_gate))
+        for circuit in circuits
+    ]
+
+    result = xeb_fitting.characterize_phased_fsim_parameters_with_xeb_by_pair(
+        sampled_df=sampled_df,
+        parameterized_circuits=p_circuits,
+        cycle_depths=cycle_depths,
+        options=options,
+        fatol=atol,
+        xatol=atol,
+        pool=pool,
+    )
+
+    return result, xeb_fitting.before_and_after_characterization(
+        fids_df_0, characterization_result=result
+    )
+
+
+def calibrate_z_phases(
+    sampler: 'cirq.Sampler',
+    qubits: Optional[Sequence['cirq.GridQubit']] = None,
+    two_qubit_gate: 'cirq.Gate' = ops.CZ,
+    options: Optional[xeb_fitting.XEBPhasedFSimCharacterizationOptions] = None,
+    n_repetitions: int = 10**4,
+    n_combinations: int = 10,
+    n_circuits: int = 20,
+    cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
+    random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    atol: float = 1e-3,
+    pool: Optional[Union[multiprocessing.Pool, concurrent.futures.ThreadPoolExecutor]] = None,
+) -> Dict[Tuple['cirq.GridQubit', 'cirq.GridQubit'], 'cirq.PhasedFSimGate']:
+    """Perform z-phase calibration for excitation-preserving gates.
+
+    For a given excitation-preserving two-qubit gate we assume an error model that can be described
+    using Z-rotations:
+                0: ───Rz(a)───two_qubit_gate───Rz(c)───
+                                │
+                1: ───Rz(b)───two_qubit_gate───Rz(d)───
+    for some angles a, b, c, and d.
+
+    Since the two-qubit gate is a excitation-preserving gate, it can be represented by an FSimGate
+    and the effect of rotations turns it into a PhasedFSimGate. Using XEB-data we find the
+    PhasedFSimGate parameters that minimize the infidelity of the gate.
+
+    References:
+        - https://arxiv.org/abs/2001.08343
+        - https://arxiv.org/abs/2010.07965
+        - https://arxiv.org/abs/1910.11333
+
+    Args:
+        sampler: The quantum engine or simulator to run the circuits.
+        qubits: Qubits to use. If none, use all qubits on the sampler's device.
+        two_qubit_gate: The entangling gate to use.
+        options: The XEB-fitting options. If None, calibrate all 5 PhasedFSimGate parameters,
+            using the representation of a two-qubit gate as an FSimGate for the initial guess.
+        n_repetitions: The number of repetitions to use.
+        n_combinations: The number of combinations to generate.
+        n_circuits: The number of circuits to generate.
+        cycle_depths: The cycle depths to use.
+        random_state: The random state to use.
+        atol: Absolute tolerance to be used by the minimizer.
+        pool: Optional multi-threading or multi-processing pool.
+
+    Returns:
+        - A dictionary mapping qubit pairs to the calibrated PhasedFSimGates.
+    """
+
+    if options is None:
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(characterize_theta=False, characterize_phi=False).with_defaults_from_gate(
+            two_qubit_gate
+        )
+
+    result, _ = z_phase_calibration_workflow(
+        sampler=sampler,
+        qubits=qubits,
+        two_qubit_gate=two_qubit_gate,
+        options=options,
+        n_repetitions=n_repetitions,
+        n_combinations=n_combinations,
+        n_circuits=n_circuits,
+        cycle_depths=cycle_depths,
+        random_state=random_state,
+        atol=atol,
+        pool=pool,
+    )
+
+    gates = {}
+    for qubits, params in result.final_params.items():
+        params['theta'] = params.get('theta', options.theta_default)
+        params['phi'] = params.get('phi', options.phi_default)
+        params['zeta'] = params.get('zeta', options.zeta_default)
+        params['chi'] = params.get('eta', options.chi_default)
+        params['gamma'] = params.get('gamma', options.gamma_default)
+        gates[qubits] = ops.PhasedFSimGate(**params)
+    return gates
\ No newline at end of file

From 145880edcf6a7b947b4a650ce6f004be5577b18f Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 2 Aug 2024 15:13:48 -0700
Subject: [PATCH 02/22] checkpoint

---
 cirq-core/cirq/experiments/__init__.py        |  6 +++++
 cirq-core/cirq/experiments/two_qubit_xeb.py   |  5 ++--
 cirq-core/cirq/experiments/xeb_fitting.py     | 12 ++++++----
 cirq-core/cirq/experiments/xeb_simulation.py  |  1 +
 .../cirq/experiments/z_phase_calibration.py   |  1 +
 .../experiments/z_phase_calibration_test.py   | 23 +++++++++++++++++++
 6 files changed, 42 insertions(+), 6 deletions(-)
 create mode 100644 cirq-core/cirq/experiments/z_phase_calibration_test.py

diff --git a/cirq-core/cirq/experiments/__init__.py b/cirq-core/cirq/experiments/__init__.py
index 71f577d4898..6efeb72331a 100644
--- a/cirq-core/cirq/experiments/__init__.py
+++ b/cirq-core/cirq/experiments/__init__.py
@@ -82,3 +82,9 @@
     parallel_two_qubit_xeb as parallel_two_qubit_xeb,
     run_rb_and_xeb as run_rb_and_xeb,
 )
+
+
+from cirq.experiments.z_phase_calibration import (
+    z_phase_calibration_workflow,
+    calibrate_z_phases,
+)
\ No newline at end of file
diff --git a/cirq-core/cirq/experiments/two_qubit_xeb.py b/cirq-core/cirq/experiments/two_qubit_xeb.py
index 5cb78c5b859..3f414e1b5b2 100644
--- a/cirq-core/cirq/experiments/two_qubit_xeb.py
+++ b/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, Union
 
 from dataclasses import dataclass
 from types import MappingProxyType
@@ -358,6 +358,7 @@ def parallel_xeb_workflow(
     cycle_depths: Sequence[int] = (5, 25, 50, 100, 200, 300),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     ax: Optional[plt.Axes] = None,
+    pool: Optional[Union['multiprocessing.Pool', 'concurrent.futures.ThreadPoolExecuter']] = None,
     **plot_kwargs,
 ) -> Tuple[pd.DataFrame, Sequence['cirq.Circuit'], pd.DataFrame]:
     """A utility method that runs the full XEB workflow.
@@ -426,7 +427,7 @@ def parallel_xeb_workflow(
     )
 
     fids = benchmark_2q_xeb_fidelities(
-        sampled_df=sampled_df, circuits=circuit_library, cycle_depths=cycle_depths
+        sampled_df=sampled_df, circuits=circuit_library, cycle_depths=cycle_depths, pool=pool
     )
 
     return fids, circuit_library, sampled_df
diff --git a/cirq-core/cirq/experiments/xeb_fitting.py b/cirq-core/cirq/experiments/xeb_fitting.py
index 55935450d67..867cc815bac 100644
--- a/cirq-core/cirq/experiments/xeb_fitting.py
+++ b/cirq-core/cirq/experiments/xeb_fitting.py
@@ -16,11 +16,13 @@
 from abc import abstractmethod, ABC
 from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
+import tqdm
 import numpy as np
 import pandas as pd
 import sympy
 from cirq import circuits, ops, protocols, _import
 from cirq.experiments.xeb_simulation import simulate_2q_xeb_circuits
+import concurrent.futures
 
 if TYPE_CHECKING:
     import cirq
@@ -41,7 +43,7 @@ def benchmark_2q_xeb_fidelities(
     circuits: Sequence['cirq.Circuit'],
     cycle_depths: Optional[Sequence[int]] = None,
     param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
-    pool: Optional['multiprocessing.pool.Pool'] = None,
+    pool: Optional[Union['multiprocessing.pool.Pool', 'concurrent.futurers.ThreadPoolExecuter']] = None,
 ) -> pd.DataFrame:
     """Simulate and benchmark two-qubit XEB circuits.
 
@@ -526,7 +528,7 @@ def characterize_phased_fsim_parameters_with_xeb_by_pair(
     initial_simplex_step_size: float = 0.1,
     xatol: float = 1e-3,
     fatol: float = 1e-3,
-    pool: Optional['multiprocessing.pool.Pool'] = None,
+    pool: Optional[Union['multiprocessing.pool.Pool', 'concurrent.futures.Executer']] = None,
 ) -> XEBCharacterizationResult:
     """Run a classical optimization to fit phased fsim parameters to experimental data, and
     thereby characterize PhasedFSim-like gates grouped by pairs.
@@ -563,11 +565,13 @@ def characterize_phased_fsim_parameters_with_xeb_by_pair(
         fatol=fatol,
     )
     subselected_dfs = [sampled_df[sampled_df['pair'] == pair] for pair in pairs]
+    # if isinstance(pool, concurrent.futures.Executor):
+    #     futures = [pool.submit(closure, df) for df in subselected_dfs]
+    #     results = [r for r in tqdm.tqdm(concurrent.futures.as_completed(futures), desc='Optimize')]
     if pool is not None:
-        results = pool.map(closure, subselected_dfs)
+        results = tqdm.tqdm(pool.map(closure, subselected_dfs), total=len(subselected_dfs), desc='Optimize Parameters')
     else:
         results = [closure(df) for df in subselected_dfs]
-
     optimization_results = {}
     all_final_params = {}
     fid_dfs = []
diff --git a/cirq-core/cirq/experiments/xeb_simulation.py b/cirq-core/cirq/experiments/xeb_simulation.py
index 25e9ce4fb1b..5c00610c30d 100644
--- a/cirq-core/cirq/experiments/xeb_simulation.py
+++ b/cirq-core/cirq/experiments/xeb_simulation.py
@@ -15,6 +15,7 @@
 from dataclasses import dataclass
 from typing import List, Optional, Sequence, TYPE_CHECKING, Dict, Any
 
+import tqdm
 import numpy as np
 import pandas as pd
 
diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index bd4052d5828..e3ae50b4651 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -87,6 +87,7 @@ def z_phase_calibration_workflow(
         n_circuits=n_circuits,
         n_combinations=n_combinations,
         random_state=random_state,
+        pool=pool,
     )
 
     if options is None:
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
new file mode 100644
index 00000000000..467da3ae006
--- /dev/null
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -0,0 +1,23 @@
+# Copyright 2024 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import cirq
+
+from cirq.experiments import z_phase_calibration_workflow, calibrate_z_phases
+
+def test_z_phase_calibration_workflow():
+    pass
+
+def test_calibrate_z_phases():
+    pass
\ No newline at end of file

From 9a9d06f70349348a9a0300429f389affdcd420cf Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 13 Sep 2024 07:52:00 -0700
Subject: [PATCH 03/22] Create workflow for Z-phase calibration

---
 cirq-core/cirq/experiments/__init__.py        |  5 +-
 cirq-core/cirq/experiments/two_qubit_xeb.py   |  6 +-
 cirq-core/cirq/experiments/xeb_fitting.py     | 27 +++---
 cirq-core/cirq/experiments/xeb_simulation.py  |  1 -
 .../cirq/experiments/z_phase_calibration.py   | 39 +++++----
 .../experiments/z_phase_calibration_test.py   | 84 +++++++++++++++++--
 6 files changed, 116 insertions(+), 46 deletions(-)

diff --git a/cirq-core/cirq/experiments/__init__.py b/cirq-core/cirq/experiments/__init__.py
index 6efeb72331a..717ea658893 100644
--- a/cirq-core/cirq/experiments/__init__.py
+++ b/cirq-core/cirq/experiments/__init__.py
@@ -84,7 +84,4 @@
 )
 
 
-from cirq.experiments.z_phase_calibration import (
-    z_phase_calibration_workflow,
-    calibrate_z_phases,
-)
\ No newline at end of file
+from cirq.experiments.z_phase_calibration import z_phase_calibration_workflow, calibrate_z_phases
diff --git a/cirq-core/cirq/experiments/two_qubit_xeb.py b/cirq-core/cirq/experiments/two_qubit_xeb.py
index 3f414e1b5b2..13bd9c78d72 100644
--- a/cirq-core/cirq/experiments/two_qubit_xeb.py
+++ b/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, Union
+from typing import Sequence, TYPE_CHECKING, Optional, Tuple, Dict, cast, Mapping
 
 from dataclasses import dataclass
 from types import MappingProxyType
@@ -38,6 +38,7 @@
 from cirq._compat import cached_method
 
 if TYPE_CHECKING:
+    import multiprocessing
     import cirq
 
 
@@ -358,7 +359,7 @@ def parallel_xeb_workflow(
     cycle_depths: Sequence[int] = (5, 25, 50, 100, 200, 300),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     ax: Optional[plt.Axes] = None,
-    pool: Optional[Union['multiprocessing.Pool', 'concurrent.futures.ThreadPoolExecuter']] = None,
+    pool: Optional['multiprocessing.pool.Pool'] = None,
     **plot_kwargs,
 ) -> Tuple[pd.DataFrame, Sequence['cirq.Circuit'], pd.DataFrame]:
     """A utility method that runs the full XEB workflow.
@@ -374,6 +375,7 @@ def parallel_xeb_workflow(
         random_state: The random state to use.
         ax: the plt.Axes to plot the device layout on. If not given,
             no plot is created.
+        pool: An optional multiprocessing pool.
         **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
 
     Returns:
diff --git a/cirq-core/cirq/experiments/xeb_fitting.py b/cirq-core/cirq/experiments/xeb_fitting.py
index 867cc815bac..c1b0d977c77 100644
--- a/cirq-core/cirq/experiments/xeb_fitting.py
+++ b/cirq-core/cirq/experiments/xeb_fitting.py
@@ -14,15 +14,13 @@
 """Estimation of fidelity associated with experimental circuit executions."""
 import dataclasses
 from abc import abstractmethod, ABC
-from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING
 
-import tqdm
 import numpy as np
 import pandas as pd
 import sympy
 from cirq import circuits, ops, protocols, _import
 from cirq.experiments.xeb_simulation import simulate_2q_xeb_circuits
-import concurrent.futures
 
 if TYPE_CHECKING:
     import cirq
@@ -43,7 +41,7 @@ def benchmark_2q_xeb_fidelities(
     circuits: Sequence['cirq.Circuit'],
     cycle_depths: Optional[Sequence[int]] = None,
     param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
-    pool: Optional[Union['multiprocessing.pool.Pool', 'concurrent.futurers.ThreadPoolExecuter']] = None,
+    pool: Optional['multiprocessing.pool.Pool'] = None,
 ) -> pd.DataFrame:
     """Simulate and benchmark two-qubit XEB circuits.
 
@@ -387,18 +385,21 @@ def SqrtISwapXEBOptions(*args, **kwargs):
 
 
 def parameterize_circuit(
-    circuit: 'cirq.Circuit', options: XEBCharacterizationOptions,     target: Union[ops.GateFamily, ops.Gateset] = ops.Gateset(
-        ops.PhasedFSimGate, ops.ISwapPowGate, ops.FSimGate
-    ),
+    circuit: 'cirq.Circuit',
+    options: XEBCharacterizationOptions,
+    target_gatefamily: Optional[ops.GateFamily] = None,
 ) -> 'cirq.Circuit':
     """Parameterize PhasedFSim-like gates in a given circuit according to
     `phased_fsim_options`.
     """
+    if isinstance(target_gatefamily, ops.GateFamily):
+        should_parameterize = lambda op: op in target_gatefamily or options.should_parameterize(op)
+    else:
+        should_parameterize = options.should_parameterize
     gate = options.get_parameterized_gate()
     return circuits.Circuit(
         circuits.Moment(
-            gate.on(*op.qubits) if op in target else op
-            for op in moment.operations
+            gate.on(*op.qubits) if should_parameterize(op) else op for op in moment.operations
         )
         for moment in circuit.moments
     )
@@ -528,7 +529,7 @@ def characterize_phased_fsim_parameters_with_xeb_by_pair(
     initial_simplex_step_size: float = 0.1,
     xatol: float = 1e-3,
     fatol: float = 1e-3,
-    pool: Optional[Union['multiprocessing.pool.Pool', 'concurrent.futures.Executer']] = None,
+    pool: Optional['multiprocessing.pool.Pool'] = None,
 ) -> XEBCharacterizationResult:
     """Run a classical optimization to fit phased fsim parameters to experimental data, and
     thereby characterize PhasedFSim-like gates grouped by pairs.
@@ -565,13 +566,11 @@ def characterize_phased_fsim_parameters_with_xeb_by_pair(
         fatol=fatol,
     )
     subselected_dfs = [sampled_df[sampled_df['pair'] == pair] for pair in pairs]
-    # if isinstance(pool, concurrent.futures.Executor):
-    #     futures = [pool.submit(closure, df) for df in subselected_dfs]
-    #     results = [r for r in tqdm.tqdm(concurrent.futures.as_completed(futures), desc='Optimize')]
     if pool is not None:
-        results = tqdm.tqdm(pool.map(closure, subselected_dfs), total=len(subselected_dfs), desc='Optimize Parameters')
+        results = pool.map(closure, subselected_dfs)
     else:
         results = [closure(df) for df in subselected_dfs]
+
     optimization_results = {}
     all_final_params = {}
     fid_dfs = []
diff --git a/cirq-core/cirq/experiments/xeb_simulation.py b/cirq-core/cirq/experiments/xeb_simulation.py
index 5c00610c30d..25e9ce4fb1b 100644
--- a/cirq-core/cirq/experiments/xeb_simulation.py
+++ b/cirq-core/cirq/experiments/xeb_simulation.py
@@ -15,7 +15,6 @@
 from dataclasses import dataclass
 from typing import List, Optional, Sequence, TYPE_CHECKING, Dict, Any
 
-import tqdm
 import numpy as np
 import pandas as pd
 
diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index e3ae50b4651..43c72fa18e1 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -13,9 +13,8 @@
 # limitations under the License.
 
 """Provides a method to do z-phase calibration for excitation-preserving gates."""
-from typing import Optional, Sequence, Union, Tuple, Dict, TYPE_CHECKING
 import multiprocessing
-import concurrent.futures
+from typing import Optional, Sequence, Tuple, Dict, TYPE_CHECKING
 
 import numpy as np
 
@@ -30,7 +29,7 @@
 
 def z_phase_calibration_workflow(
     sampler: 'cirq.Sampler',
-    qubits: Optional['cirq.GridQubit'] = None,
+    qubits: Optional[Sequence['cirq.GridQubit']] = None,
     two_qubit_gate: 'cirq.Gate' = ops.CZ,
     options: Optional[xeb_fitting.XEBPhasedFSimCharacterizationOptions] = None,
     n_repetitions: int = 10**4,
@@ -39,7 +38,7 @@ def z_phase_calibration_workflow(
     cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
-    pool: Optional[Union[multiprocessing.Pool, concurrent.futures.ThreadPoolExecutor]] = None,
+    pool: Optional[multiprocessing.pool.Pool] = None,
 ) -> Tuple[xeb_fitting.XEBCharacterizationResult, 'pd.DataFrame']:
     """Perform z-phase calibration for excitation-preserving gates.
 
@@ -91,9 +90,9 @@ def z_phase_calibration_workflow(
     )
 
     if options is None:
-        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(characterize_theta=False, characterize_phi=False).with_defaults_from_gate(
-            two_qubit_gate
-        )
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
+            characterize_theta=False, characterize_phi=False
+        ).with_defaults_from_gate(two_qubit_gate)
 
     p_circuits = [
         xeb_fitting.parameterize_circuit(circuit, options, ops.GateFamily(two_qubit_gate))
@@ -126,8 +125,8 @@ def calibrate_z_phases(
     cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
-    pool: Optional[Union[multiprocessing.Pool, concurrent.futures.ThreadPoolExecutor]] = None,
-) -> Dict[Tuple['cirq.GridQubit', 'cirq.GridQubit'], 'cirq.PhasedFSimGate']:
+    pool: Optional[multiprocessing.pool.Pool] = None,
+) -> Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate']:
     """Perform z-phase calibration for excitation-preserving gates.
 
     For a given excitation-preserving two-qubit gate we assume an error model that can be described
@@ -165,9 +164,9 @@ def calibrate_z_phases(
     """
 
     if options is None:
-        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(characterize_theta=False, characterize_phi=False).with_defaults_from_gate(
-            two_qubit_gate
-        )
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
+            characterize_theta=False, characterize_phi=False
+        ).with_defaults_from_gate(two_qubit_gate)
 
     result, _ = z_phase_calibration_workflow(
         sampler=sampler,
@@ -184,11 +183,11 @@ def calibrate_z_phases(
     )
 
     gates = {}
-    for qubits, params in result.final_params.items():
-        params['theta'] = params.get('theta', options.theta_default)
-        params['phi'] = params.get('phi', options.phi_default)
-        params['zeta'] = params.get('zeta', options.zeta_default)
-        params['chi'] = params.get('eta', options.chi_default)
-        params['gamma'] = params.get('gamma', options.gamma_default)
-        gates[qubits] = ops.PhasedFSimGate(**params)
-    return gates
\ No newline at end of file
+    for pair, params in result.final_params.items():
+        params['theta'] = params.get('theta', options.theta_default or 0)
+        params['phi'] = params.get('phi', options.phi_default or 0)
+        params['zeta'] = params.get('zeta', options.zeta_default or 0)
+        params['chi'] = params.get('eta', options.chi_default or 0)
+        params['gamma'] = params.get('gamma', options.gamma_default or 0)
+        gates[pair] = ops.PhasedFSimGate(**params)
+    return gates
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 467da3ae006..b0733f458a0 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -12,12 +12,86 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
+import numpy as np
+
 import cirq
 
-from cirq.experiments import z_phase_calibration_workflow, calibrate_z_phases
+from cirq.experiments import calibrate_z_phases
+from cirq.experiments.xeb_fitting import XEBPhasedFSimCharacterizationOptions
+
+_ANGLES = ['theta', 'phi', 'chi', 'zeta', 'gamma']
+
+
+def _create_tests(n, seed):
+    rng = np.random.default_rng(seed)
+    angles = (rng.random((n, 5)) * 2 - 1) * np.pi
+    # Add errors to the first 2 angles (theta and phi).
+    # The errors for theta and phi are in the interval (-1, 1).
+    error = np.concatenate(
+        [rng.random((n, 2)) * rng.choice([-1, 1], (n, 2)), np.zeros((n, 3))], axis=-1
+    )
+    return zip(angles, error)
+
+
+def _trace_distance(A, B):
+    return 0.5 * np.abs(np.linalg.eigvals(A - B)).sum()
+
+
+class _TestSimulator(cirq.Simulator):
+    """A simulator that replaces a specific gate by another."""
+
+    def __init__(self, gate: cirq.Gate, replacement: cirq.Gate, **kwargs):
+        super().__init__(**kwargs)
+        self.gate = gate
+        self.replacement = replacement
+
+    def _core_iterator(
+        self,
+        circuit: 'cirq.AbstractCircuit',
+        sim_state,
+        all_measurements_are_terminal: bool = False,
+    ):
+        new_circuit = cirq.Circuit(
+            [
+                [op if op.gate != self.gate else self.replacement(*op.qubits) for op in m]
+                for m in circuit
+            ]
+        )
+        yield from super()._core_iterator(new_circuit, sim_state, all_measurements_are_terminal)
+
+
+@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=10, seed=0))
+def test_calibrate_z_phases(angles, error):
+
+    original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})
+    actual_gate = cirq.PhasedFSimGate(**{k: v + e for k, v, e in zip(_ANGLES, angles, error)})
+
+    options = XEBPhasedFSimCharacterizationOptions(
+        **{f'{n}_default': t for n, t in zip(_ANGLES, angles)},
+        characterize_chi=False,
+        characterize_gamma=False,
+        characterize_phi=True,
+        characterize_theta=True,
+        characterize_zeta=False,
+    )
 
-def test_z_phase_calibration_workflow():
-    pass
+    sampler = _TestSimulator(original_gate, actual_gate, seed=0)
+    qubits = cirq.q(0, 0), cirq.q(0, 1)
+    calibrated_gate = calibrate_z_phases(
+        sampler,
+        qubits,
+        original_gate,
+        options,
+        n_repetitions=10,
+        n_combinations=10,
+        n_circuits=10,
+        cycle_depths=range(3, 10),
+    )[qubits]
 
-def test_calibrate_z_phases():
-    pass
\ No newline at end of file
+    initial_unitary = cirq.unitary(original_gate)
+    final_unitary = cirq.unitary(calibrated_gate)
+    target_unitary = cirq.unitary(actual_gate)
+    assert _trace_distance(final_unitary, target_unitary) < _trace_distance(
+        initial_unitary, target_unitary
+    )

From d37ea6068c9ead203e63744f8f8ab0c851ba59c6 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 13 Sep 2024 08:02:14 -0700
Subject: [PATCH 04/22] update deps

---
 cirq-core/cirq/experiments/z_phase_calibration.py      | 6 +++---
 cirq-core/cirq/experiments/z_phase_calibration_test.py | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 43c72fa18e1..dcc7d6ca314 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 """Provides a method to do z-phase calibration for excitation-preserving gates."""
-import multiprocessing
 from typing import Optional, Sequence, Tuple, Dict, TYPE_CHECKING
 
 import numpy as np
@@ -25,6 +24,7 @@
 if TYPE_CHECKING:
     import cirq
     import pandas as pd
+    import multiprocessing
 
 
 def z_phase_calibration_workflow(
@@ -38,7 +38,7 @@ def z_phase_calibration_workflow(
     cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
-    pool: Optional[multiprocessing.pool.Pool] = None,
+    pool: Optional['multiprocessing.pool.Pool'] = None,
 ) -> Tuple[xeb_fitting.XEBCharacterizationResult, 'pd.DataFrame']:
     """Perform z-phase calibration for excitation-preserving gates.
 
@@ -125,7 +125,7 @@ def calibrate_z_phases(
     cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
-    pool: Optional[multiprocessing.pool.Pool] = None,
+    pool: Optional['multiprocessing.pool.Pool'] = None,
 ) -> Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate']:
     """Perform z-phase calibration for excitation-preserving gates.
 
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index b0733f458a0..a02e60635b9 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -17,7 +17,7 @@
 
 import cirq
 
-from cirq.experiments import calibrate_z_phases
+from cirq.experiments.z_phase_calibration import calibrate_z_phases
 from cirq.experiments.xeb_fitting import XEBPhasedFSimCharacterizationOptions
 
 _ANGLES = ['theta', 'phi', 'chi', 'zeta', 'gamma']

From 717155c3f1656013954050d16d50e87ded76d6d9 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 13 Sep 2024 08:14:12 -0700
Subject: [PATCH 05/22] increase error

---
 cirq-core/cirq/experiments/z_phase_calibration_test.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index a02e60635b9..a9ff4823bfd 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -27,9 +27,10 @@ def _create_tests(n, seed):
     rng = np.random.default_rng(seed)
     angles = (rng.random((n, 5)) * 2 - 1) * np.pi
     # Add errors to the first 2 angles (theta and phi).
-    # The errors for theta and phi are in the interval (-1, 1).
+    # The errors for theta and phi are in the union (-1, -0.5) U (0.5, 1).
+    # This is because we run the tests with few repetitions so a small error might not get fixed.
     error = np.concatenate(
-        [rng.random((n, 2)) * rng.choice([-1, 1], (n, 2)), np.zeros((n, 3))], axis=-1
+        [(rng.random((n, 2)) * 0.5 + 0.5) * rng.choice([-1, 1], (n, 2)), np.zeros((n, 3))], axis=-1
     )
     return zip(angles, error)
 

From b053a9361a2314080a0dc17045a25bfd66882e23 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 13 Sep 2024 08:26:58 -0700
Subject: [PATCH 06/22] nit

---
 cirq-core/cirq/experiments/z_phase_calibration_test.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index a9ff4823bfd..a4d20746ebf 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -27,10 +27,10 @@ def _create_tests(n, seed):
     rng = np.random.default_rng(seed)
     angles = (rng.random((n, 5)) * 2 - 1) * np.pi
     # Add errors to the first 2 angles (theta and phi).
-    # The errors for theta and phi are in the union (-1, -0.5) U (0.5, 1).
+    # The errors for theta and phi are in the union (-2, -1) U (1, 2).
     # This is because we run the tests with few repetitions so a small error might not get fixed.
     error = np.concatenate(
-        [(rng.random((n, 2)) * 0.5 + 0.5) * rng.choice([-1, 1], (n, 2)), np.zeros((n, 3))], axis=-1
+        [(rng.random((n, 2)) + 1) * rng.choice([-1, 1], (n, 2)), np.zeros((n, 3))], axis=-1
     )
     return zip(angles, error)
 
@@ -62,7 +62,7 @@ def _core_iterator(
         yield from super()._core_iterator(new_circuit, sim_state, all_measurements_are_terminal)
 
 
-@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=10, seed=0))
+@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=10, seed=32432432))
 def test_calibrate_z_phases(angles, error):
 
     original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})

From eadda56ba60738c79fcf313a053d68978dfbd110 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 13 Sep 2024 08:52:47 -0700
Subject: [PATCH 07/22] nit

---
 .../cirq/experiments/z_phase_calibration_test.py    | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index a4d20746ebf..7405087cdf6 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -93,6 +93,13 @@ def test_calibrate_z_phases(angles, error):
     initial_unitary = cirq.unitary(original_gate)
     final_unitary = cirq.unitary(calibrated_gate)
     target_unitary = cirq.unitary(actual_gate)
-    assert _trace_distance(final_unitary, target_unitary) < _trace_distance(
-        initial_unitary, target_unitary
-    )
+    maximally_mixed_state = np.eye(4) / 2
+    dm_initial = initial_unitary @ maximally_mixed_state @ initial_unitary.T.conj()
+    dm_final = final_unitary @ maximally_mixed_state @ final_unitary.T.conj()
+    dm_target = target_unitary @ maximally_mixed_state @ target_unitary.T.conj()
+
+    original_dist = _trace_distance(dm_initial, dm_target)
+    new_dist = _trace_distance(dm_final, dm_target)
+
+    # Either we reduced the error or the error is small enough.
+    assert new_dist < original_dist or new_dist < 1e-6

From 07ba30a0895df8b13697f900972a8d7f51caaf7c Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Fri, 13 Sep 2024 08:59:32 -0700
Subject: [PATCH 08/22] Add test

---
 .../cirq/experiments/z_phase_calibration.py   | 12 +++----
 .../experiments/z_phase_calibration_test.py   | 35 +++++++++++++++++++
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index dcc7d6ca314..41d168e1fb6 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -90,9 +90,9 @@ def z_phase_calibration_workflow(
     )
 
     if options is None:
-        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
-            characterize_theta=False, characterize_phi=False
-        ).with_defaults_from_gate(two_qubit_gate)
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(
+            two_qubit_gate
+        )
 
     p_circuits = [
         xeb_fitting.parameterize_circuit(circuit, options, ops.GateFamily(two_qubit_gate))
@@ -164,9 +164,9 @@ def calibrate_z_phases(
     """
 
     if options is None:
-        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
-            characterize_theta=False, characterize_phi=False
-        ).with_defaults_from_gate(two_qubit_gate)
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(
+            two_qubit_gate
+        )
 
     result, _ = z_phase_calibration_workflow(
         sampler=sampler,
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 7405087cdf6..648c5d8a9c8 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -103,3 +103,38 @@ def test_calibrate_z_phases(angles, error):
 
     # Either we reduced the error or the error is small enough.
     assert new_dist < original_dist or new_dist < 1e-6
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=10, seed=32432432))
+def test_calibrate_z_phases_no_options(angles, error):
+
+    original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})
+    actual_gate = cirq.PhasedFSimGate(**{k: v + e for k, v, e in zip(_ANGLES, angles, error)})
+
+    sampler = _TestSimulator(original_gate, actual_gate, seed=0)
+    qubits = cirq.q(0, 0), cirq.q(0, 1)
+    calibrated_gate = calibrate_z_phases(
+        sampler,
+        qubits,
+        original_gate,
+        options=None,
+        n_repetitions=10,
+        n_combinations=10,
+        n_circuits=10,
+        cycle_depths=range(3, 10),
+    )[qubits]
+
+    initial_unitary = cirq.unitary(original_gate)
+    final_unitary = cirq.unitary(calibrated_gate)
+    target_unitary = cirq.unitary(actual_gate)
+    maximally_mixed_state = np.eye(4) / 2
+    dm_initial = initial_unitary @ maximally_mixed_state @ initial_unitary.T.conj()
+    dm_final = final_unitary @ maximally_mixed_state @ final_unitary.T.conj()
+    dm_target = target_unitary @ maximally_mixed_state @ target_unitary.T.conj()
+
+    original_dist = _trace_distance(dm_initial, dm_target)
+    new_dist = _trace_distance(dm_final, dm_target)
+
+    # Either we reduced the error or the error is small enough.
+    assert new_dist < original_dist or new_dist < 1e-6

From 552b7dfa6a52b929fa3259e9e1a7fcac0848b032 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 18 Sep 2024 23:32:20 -0700
Subject: [PATCH 09/22] checkpoint

---
 cirq-core/cirq/experiments/z_phase_calibration.py      | 4 ++--
 cirq-core/cirq/experiments/z_phase_calibration_test.py | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 41d168e1fb6..d16165297b8 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -90,7 +90,7 @@ def z_phase_calibration_workflow(
     )
 
     if options is None:
-        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(characterize_chi=False, characterize_gamma=False, characterize_zeta=False).with_defaults_from_gate(
             two_qubit_gate
         )
 
@@ -164,7 +164,7 @@ def calibrate_z_phases(
     """
 
     if options is None:
-        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions().with_defaults_from_gate(
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(characterize_chi=False, characterize_gamma=False, characterize_zeta=False).with_defaults_from_gate(
             two_qubit_gate
         )
 
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 648c5d8a9c8..9458aa4dab7 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -105,7 +105,6 @@ def test_calibrate_z_phases(angles, error):
     assert new_dist < original_dist or new_dist < 1e-6
 
 
-@pytest.mark.slow
 @pytest.mark.parametrize(['angles', 'error'], _create_tests(n=10, seed=32432432))
 def test_calibrate_z_phases_no_options(angles, error):
 

From 6f71abf2cce2e40524e0964dd9c3aaa7ecd6e2c2 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 18 Sep 2024 23:48:39 -0700
Subject: [PATCH 10/22] coverage

---
 .../cirq/experiments/z_phase_calibration.py   | 14 +++++-----
 .../experiments/z_phase_calibration_test.py   | 27 ++++++++++++++++++-
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index d16165297b8..058afac81b5 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -90,9 +90,9 @@ def z_phase_calibration_workflow(
     )
 
     if options is None:
-        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(characterize_chi=False, characterize_gamma=False, characterize_zeta=False).with_defaults_from_gate(
-            two_qubit_gate
-        )
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
+            characterize_chi=False, characterize_gamma=False, characterize_zeta=False
+        ).with_defaults_from_gate(two_qubit_gate)
 
     p_circuits = [
         xeb_fitting.parameterize_circuit(circuit, options, ops.GateFamily(two_qubit_gate))
@@ -164,9 +164,9 @@ def calibrate_z_phases(
     """
 
     if options is None:
-        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(characterize_chi=False, characterize_gamma=False, characterize_zeta=False).with_defaults_from_gate(
-            two_qubit_gate
-        )
+        options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
+            characterize_chi=False, characterize_gamma=False, characterize_zeta=False
+        ).with_defaults_from_gate(two_qubit_gate)
 
     result, _ = z_phase_calibration_workflow(
         sampler=sampler,
@@ -187,7 +187,7 @@ def calibrate_z_phases(
         params['theta'] = params.get('theta', options.theta_default or 0)
         params['phi'] = params.get('phi', options.phi_default or 0)
         params['zeta'] = params.get('zeta', options.zeta_default or 0)
-        params['chi'] = params.get('eta', options.chi_default or 0)
+        params['chi'] = params.get('chi', options.chi_default or 0)
         params['gamma'] = params.get('gamma', options.gamma_default or 0)
         gates[pair] = ops.PhasedFSimGate(**params)
     return gates
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 9458aa4dab7..581b65ee68a 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -17,7 +17,7 @@
 
 import cirq
 
-from cirq.experiments.z_phase_calibration import calibrate_z_phases
+from cirq.experiments.z_phase_calibration import calibrate_z_phases, z_phase_calibration_workflow
 from cirq.experiments.xeb_fitting import XEBPhasedFSimCharacterizationOptions
 
 _ANGLES = ['theta', 'phi', 'chi', 'zeta', 'gamma']
@@ -137,3 +137,28 @@ def test_calibrate_z_phases_no_options(angles, error):
 
     # Either we reduced the error or the error is small enough.
     assert new_dist < original_dist or new_dist < 1e-6
+
+
+@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=10, seed=32432432))
+def test_calibrate_z_phases_workflow_no_options(angles, error):
+
+    original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})
+    actual_gate = cirq.PhasedFSimGate(**{k: v + e for k, v, e in zip(_ANGLES, angles, error)})
+
+    sampler = _TestSimulator(original_gate, actual_gate, seed=0)
+    qubits = cirq.q(0, 0), cirq.q(0, 1)
+    result, _ = z_phase_calibration_workflow(
+        sampler,
+        qubits,
+        original_gate,
+        options=None,
+        n_repetitions=1,
+        n_combinations=1,
+        n_circuits=1,
+        cycle_depths=(1, 2),
+    )
+
+    for params in result.final_params.values():
+        assert 'zeta' not in params
+        assert 'chi' not in params
+        assert 'gamma' not in params

From 39b1537f34781e920963bb538bbb81c84781f1b9 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 2 Oct 2024 13:29:12 -0700
Subject: [PATCH 11/22] Add plotting method and calculate the variance of
 estimated fidelity

---
 cirq-core/cirq/experiments/xeb_fitting.py     | 13 +++++-
 .../cirq/experiments/z_phase_calibration.py   | 41 ++++++++++++++++++-
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/cirq-core/cirq/experiments/xeb_fitting.py b/cirq-core/cirq/experiments/xeb_fitting.py
index c1b0d977c77..6e797dc8c2e 100644
--- a/cirq-core/cirq/experiments/xeb_fitting.py
+++ b/cirq-core/cirq/experiments/xeb_fitting.py
@@ -95,6 +95,11 @@ def benchmark_2q_xeb_fidelities(
     df['e_u'] = np.sum(pure_probs**2, axis=1)
     df['u_u'] = np.sum(pure_probs, axis=1) / D
     df['m_u'] = np.sum(pure_probs * sampled_probs, axis=1)
+    # Var[m_u] = Var[sum p(x) * p_sampled(x)]
+    #           = sum p(x)^2 Var[p_sampled(x)]
+    #           = sum p(x)^2 p(x) (1 - p(x))
+    #           = sum p(x)^3 (1 - p(x))
+    df['var_m_u'] = np.sum(pure_probs**3 * (1 - pure_probs), axis=1)
     df['y'] = df['m_u'] - df['u_u']
     df['x'] = df['e_u'] - df['u_u']
     df['numerator'] = df['x'] * df['y']
@@ -103,7 +108,11 @@ def benchmark_2q_xeb_fidelities(
     def per_cycle_depth(df):
         """This function is applied per cycle_depth in the following groupby aggregation."""
         fid_lsq = df['numerator'].sum() / df['denominator'].sum()
-        ret = {'fidelity': fid_lsq}
+        # Note: both df['denominator'] an df['x'] are constants.
+        # Var[f] = Var[df['numerator']] / (sum df['denominator'])^2
+        #           = sum (df['x']^2 * df['var_m_u']) / (sum df['denominator'])^2
+        var_fid = (df['var_m_u'] * df['x'] ** 2).sum() / df['denominator'].sum() ** 2
+        ret = {'fidelity': fid_lsq, 'fidelity_variance': var_fid}
 
         def _try_keep(k):
             """If all the values for a key `k` are the same in this group, we can keep it."""
@@ -678,7 +687,7 @@ def _per_pair(f1):
             'cycle_depths': f1['cycle_depth'].values,
             'fidelities': f1['fidelity'].values,
             'a_std': a_std,
-            'layer_fid_std': layer_fid_std,
+            'layer_fid_std': np.sqrt(layer_fid_std**2 + f1['fidelity_variance'].values),
         }
         return pd.Series(record)
 
diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 058afac81b5..0f15b752744 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -25,6 +25,7 @@
     import cirq
     import pandas as pd
     import multiprocessing
+    import matplotlib.pyplot as plt
 
 
 def z_phase_calibration_workflow(
@@ -109,10 +110,12 @@ def z_phase_calibration_workflow(
         pool=pool,
     )
 
-    return result, xeb_fitting.before_and_after_characterization(
+    before_after = xeb_fitting.before_and_after_characterization(
         fids_df_0, characterization_result=result
     )
 
+    return result, before_after
+
 
 def calibrate_z_phases(
     sampler: 'cirq.Sampler',
@@ -191,3 +194,39 @@ def calibrate_z_phases(
         params['gamma'] = params.get('gamma', options.gamma_default or 0)
         gates[pair] = ops.PhasedFSimGate(**params)
     return gates
+
+
+def plot_z_phase_calibration_result(
+    before_after_df: 'pd.DataFrame',
+    axes: np.ndarray[Sequence[Sequence['plt.Axes']], np.dtype[np.object_]],
+    *,
+    with_error_bars: bool = False,
+) -> None:
+    """A helper method to plot the result of running z-phase calibration.
+
+    Args:
+        before_after_df: The second return object of running `z_phase_calibration_workflow`.
+        axes: And ndarray of the axes to plot on.
+            The number of axes is expected to be >= number of qubit pairs.
+        with_error_bars: Whether to add error bars or not.
+            The width of the bar is an upper bound on standard variation of the estimated fidelity.
+    """
+    for pair, ax in zip(before_after_df.index, axes.flatten()):
+        row = before_after_df.loc[[pair]].iloc[0]
+        ax.errorbar(
+            row.cycle_depths_0,
+            row.fidelities_0,
+            yerr=row.layer_fid_std_0 * with_error_bars,
+            label='original',
+        )
+        ax.errorbar(
+            row.cycle_depths_0,
+            row.fidelities_c,
+            yerr=row.layer_fid_std_c * with_error_bars,
+            label='calibrated',
+        )
+        ax.axhline(1, linestyle='--')
+        ax.set_xlabel('cycle depth')
+        ax.set_ylabel('fidelity estimate')
+        ax.set_title('-'.join(str(q) for q in pair))
+        ax.legend()

From 9f42000f7f340a6cc9c12035ffd515f8e086cefe Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 2 Oct 2024 13:33:35 -0700
Subject: [PATCH 12/22] docstring

---
 cirq-core/cirq/experiments/z_phase_calibration.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 0f15b752744..f12e76bdcb4 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -204,6 +204,9 @@ def plot_z_phase_calibration_result(
 ) -> None:
     """A helper method to plot the result of running z-phase calibration.
 
+    Note that the plotted fidelity is a statistical estimate of the true fidelity and as a result
+    may be outside the [0, 1] range.
+
     Args:
         before_after_df: The second return object of running `z_phase_calibration_workflow`.
         axes: And ndarray of the axes to plot on.

From 26593a17bd5d0f4359a8ee2012d8dd75808479f4 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 2 Oct 2024 14:39:11 -0700
Subject: [PATCH 13/22] coverage

---
 cirq-core/cirq/experiments/xeb_fitting.py     |  5 +++-
 .../experiments/z_phase_calibration_test.py   | 29 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/cirq-core/cirq/experiments/xeb_fitting.py b/cirq-core/cirq/experiments/xeb_fitting.py
index 6e797dc8c2e..ae16e54d478 100644
--- a/cirq-core/cirq/experiments/xeb_fitting.py
+++ b/cirq-core/cirq/experiments/xeb_fitting.py
@@ -681,13 +681,16 @@ def _per_pair(f1):
         a, layer_fid, a_std, layer_fid_std = _fit_exponential_decay(
             f1['cycle_depth'], f1['fidelity']
         )
+        fidelity_variance = 0
+        if 'fidelity_variance' in f1:
+            fidelity_variance = f1['fidelity_variance'].values
         record = {
             'a': a,
             'layer_fid': layer_fid,
             'cycle_depths': f1['cycle_depth'].values,
             'fidelities': f1['fidelity'].values,
             'a_std': a_std,
-            'layer_fid_std': np.sqrt(layer_fid_std**2 + f1['fidelity_variance'].values),
+            'layer_fid_std': np.sqrt(layer_fid_std**2 + fidelity_variance),
         }
         return pd.Series(record)
 
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 581b65ee68a..b766b79c9bd 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -14,10 +14,16 @@
 
 import pytest
 import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
 
 import cirq
 
-from cirq.experiments.z_phase_calibration import calibrate_z_phases, z_phase_calibration_workflow
+from cirq.experiments.z_phase_calibration import (
+    calibrate_z_phases,
+    z_phase_calibration_workflow,
+    plot_z_phase_calibration_result,
+)
 from cirq.experiments.xeb_fitting import XEBPhasedFSimCharacterizationOptions
 
 _ANGLES = ['theta', 'phi', 'chi', 'zeta', 'gamma']
@@ -162,3 +168,24 @@ def test_calibrate_z_phases_workflow_no_options(angles, error):
         assert 'zeta' not in params
         assert 'chi' not in params
         assert 'gamma' not in params
+
+
+def test_plot_z_phase_calibration_result():
+    df = pd.DataFrame()
+    df.index = [(1, 2), (2, 3)]
+    df['cycle_depths_0'] = [[1, 2, 3]] * 2
+    df['fidelities_0'] = [[0.9, 0.8, 0.7], [0.6, 0.4, 0.1]]
+    df['layer_fid_std_0'] = [0.1, 0.2]
+    df['fidelities_c'] = [[0.9, 0.92, 0.93], [0.7, 0.77, 0.8]]
+    df['layer_fid_std_c'] = [0.2, 0.3]
+
+    _, axes = plt.subplots(1, 2)
+    plot_z_phase_calibration_result(before_after_df=df, axes=axes)
+
+    np.testing.assert_allclose(axes[0].lines[0].get_xdata().astype(float), [1, 2, 3])
+    np.testing.assert_allclose(axes[0].lines[0].get_ydata().astype(float), [0.9, 0.8, 0.7])
+    np.testing.assert_allclose(axes[0].lines[1].get_ydata().astype(float), [0.9, 0.92, 0.93])
+
+    np.testing.assert_allclose(axes[1].lines[0].get_xdata().astype(float), [1, 2, 3])
+    np.testing.assert_allclose(axes[1].lines[0].get_ydata().astype(float), [0.6, 0.4, 0.1])
+    np.testing.assert_allclose(axes[1].lines[1].get_ydata().astype(float), [0.7, 0.77, 0.8])

From 39e0a6d0722506fc90b8fc84b51e42958e79f7e0 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 2 Oct 2024 15:56:52 -0700
Subject: [PATCH 14/22] change qubit names

---
 cirq-core/cirq/experiments/z_phase_calibration_test.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index b766b79c9bd..0d92b6756ee 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -172,7 +172,8 @@ def test_calibrate_z_phases_workflow_no_options(angles, error):
 
 def test_plot_z_phase_calibration_result():
     df = pd.DataFrame()
-    df.index = [(1, 2), (2, 3)]
+    qs = cirq.q(0, 0), cirq.q(0, 1), cirq.q(0, 2)
+    df.index = [qs[:2], qs[-2:]]
     df['cycle_depths_0'] = [[1, 2, 3]] * 2
     df['fidelities_0'] = [[0.9, 0.8, 0.7], [0.6, 0.4, 0.1]]
     df['layer_fid_std_0'] = [0.1, 0.2]

From 530dc2dfff526e02c5d3be575896d6226c9d48dc Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 30 Oct 2024 16:13:59 -0700
Subject: [PATCH 15/22] checkpoint

---
 cirq-core/cirq/experiments/xeb_fitting.py     |  2 +-
 .../cirq/experiments/z_phase_calibration.py   | 63 ++++++++++++++-----
 .../experiments/z_phase_calibration_test.py   | 48 +++++++++-----
 3 files changed, 82 insertions(+), 31 deletions(-)

diff --git a/cirq-core/cirq/experiments/xeb_fitting.py b/cirq-core/cirq/experiments/xeb_fitting.py
index ae16e54d478..54e9dfcd14c 100644
--- a/cirq-core/cirq/experiments/xeb_fitting.py
+++ b/cirq-core/cirq/experiments/xeb_fitting.py
@@ -108,7 +108,7 @@ def benchmark_2q_xeb_fidelities(
     def per_cycle_depth(df):
         """This function is applied per cycle_depth in the following groupby aggregation."""
         fid_lsq = df['numerator'].sum() / df['denominator'].sum()
-        # Note: both df['denominator'] an df['x'] are constants.
+        # Note: both df['denominator'] and df['x'] are constants.
         # Var[f] = Var[df['numerator']] / (sum df['denominator'])^2
         #           = sum (df['x']^2 * df['var_m_u']) / (sum df['denominator'])^2
         var_fid = (df['var_m_u'] * df['x'] ** 2).sum() / df['denominator'].sum() ** 2
diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index f12e76bdcb4..6638689a162 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -13,8 +13,11 @@
 # limitations under the License.
 
 """Provides a method to do z-phase calibration for excitation-preserving gates."""
-from typing import Optional, Sequence, Tuple, Dict, TYPE_CHECKING
+from typing import Callable, Union, Optional, Sequence, Tuple, Dict, TYPE_CHECKING
 
+import multiprocessing
+import multiprocessing.pool
+import matplotlib.pyplot as plt
 import numpy as np
 
 from cirq.experiments import xeb_fitting
@@ -24,8 +27,6 @@
 if TYPE_CHECKING:
     import cirq
     import pandas as pd
-    import multiprocessing
-    import matplotlib.pyplot as plt
 
 
 def z_phase_calibration_workflow(
@@ -39,7 +40,7 @@ def z_phase_calibration_workflow(
     cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
-    pool: Optional['multiprocessing.pool.Pool'] = None,
+    pool_or_num_workers: Optional[Union[int, 'multiprocessing.pool.Pool']] = None,
 ) -> Tuple[xeb_fitting.XEBCharacterizationResult, 'pd.DataFrame']:
     """Perform z-phase calibration for excitation-preserving gates.
 
@@ -71,13 +72,23 @@ def z_phase_calibration_workflow(
         cycle_depths: The cycle depths to use.
         random_state: The random state to use.
         atol: Absolute tolerance to be used by the minimizer.
-        pool: Optional multi-threading or multi-processing pool.
-
+        pool_or_num_workers: An optional multi-processing pool or number of workers.
+            A zero value means no multiprocessing.
+            A positivie integer value will create a pool with the given number of workers.
+            A None value will create pool with maximum number of workers.
     Returns:
         - An `XEBCharacterizationResult` object that contains the calibration result.
-        - A `pd.DataFrame` comparing the before and after fidilities.
+        - A `pd.DataFrame` comparing the before and after fidelities.
     """
 
+    pool: Optional['multiprocessing.pool.Pool'] = None
+    local_pool = False
+    if isinstance(pool_or_num_workers, multiprocessing.pool.Pool):
+        pool = pool_or_num_workers  # pragma: no cover
+    elif pool_or_num_workers != 0:
+        pool = multiprocessing.Pool(pool_or_num_workers)
+        local_pool = True
+
     fids_df_0, circuits, sampled_df = parallel_xeb_workflow(
         sampler=sampler,
         qubits=qubits,
@@ -92,7 +103,11 @@ def z_phase_calibration_workflow(
 
     if options is None:
         options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
-            characterize_chi=False, characterize_gamma=False, characterize_zeta=False
+            characterize_chi=True,
+            characterize_gamma=True,
+            characterize_zeta=True,
+            characterize_theta=False,
+            characterize_phi=False,
         ).with_defaults_from_gate(two_qubit_gate)
 
     p_circuits = [
@@ -114,6 +129,8 @@ def z_phase_calibration_workflow(
         fids_df_0, characterization_result=result
     )
 
+    if local_pool:
+        pool.close()
     return result, before_after
 
 
@@ -128,7 +145,7 @@ def calibrate_z_phases(
     cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
-    pool: Optional['multiprocessing.pool.Pool'] = None,
+    pool_or_num_workers: Optional[Union[int, 'multiprocessing.pool.Pool']] = None,
 ) -> Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate']:
     """Perform z-phase calibration for excitation-preserving gates.
 
@@ -160,7 +177,10 @@ def calibrate_z_phases(
         cycle_depths: The cycle depths to use.
         random_state: The random state to use.
         atol: Absolute tolerance to be used by the minimizer.
-        pool: Optional multi-threading or multi-processing pool.
+        pool_or_num_workers: An optional multi-processing pool or number of workers.
+            A zero value means no multiprocessing.
+            A positivie integer value will create a pool with the given number of workers.
+            A None value will create pool with maximum number of workers.
 
     Returns:
         - A dictionary mapping qubit pairs to the calibrated PhasedFSimGates.
@@ -168,7 +188,11 @@ def calibrate_z_phases(
 
     if options is None:
         options = xeb_fitting.XEBPhasedFSimCharacterizationOptions(
-            characterize_chi=False, characterize_gamma=False, characterize_zeta=False
+            characterize_chi=True,
+            characterize_gamma=True,
+            characterize_zeta=True,
+            characterize_theta=False,
+            characterize_phi=False,
         ).with_defaults_from_gate(two_qubit_gate)
 
     result, _ = z_phase_calibration_workflow(
@@ -182,7 +206,7 @@ def calibrate_z_phases(
         cycle_depths=cycle_depths,
         random_state=random_state,
         atol=atol,
-        pool=pool,
+        pool_or_num_workers=pool_or_num_workers,
     )
 
     gates = {}
@@ -198,7 +222,8 @@ def calibrate_z_phases(
 
 def plot_z_phase_calibration_result(
     before_after_df: 'pd.DataFrame',
-    axes: np.ndarray[Sequence[Sequence['plt.Axes']], np.dtype[np.object_]],
+    axes: Optional[Sequence['plt.Axes']] = None,
+    pairs: Optional[Sequence[Tuple['cirq.Qid', 'cirq.Qid']]] = None,
     *,
     with_error_bars: bool = False,
 ) -> None:
@@ -211,10 +236,19 @@ def plot_z_phase_calibration_result(
         before_after_df: The second return object of running `z_phase_calibration_workflow`.
         axes: And ndarray of the axes to plot on.
             The number of axes is expected to be >= number of qubit pairs.
+        pairs: If provided, only the given pairs are plotted.
         with_error_bars: Whether to add error bars or not.
             The width of the bar is an upper bound on standard variation of the estimated fidelity.
     """
-    for pair, ax in zip(before_after_df.index, axes.flatten()):
+    if pairs is None:
+        pairs = before_after_df.index
+    if axes is None:
+        # Create a 16x9 rectangle.
+        ncols = int(np.ceil(np.sqrt(9 / 16 * len(pairs))))
+        nrows = (len(pairs) + ncols - 1) // ncols
+        _, axes = plt.subplots(nrows=nrows, ncols=ncols)
+    axes = axes if isinstance(axes, np.ndarray) else np.array(axes)
+    for pair, ax in zip(pairs, axes.flatten()):
         row = before_after_df.loc[[pair]].iloc[0]
         ax.errorbar(
             row.cycle_depths_0,
@@ -233,3 +267,4 @@ def plot_z_phase_calibration_result(
         ax.set_ylabel('fidelity estimate')
         ax.set_title('-'.join(str(q) for q in pair))
         ax.legend()
+    return axes
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 0d92b6756ee..2e29d48e8c8 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -29,7 +29,7 @@
 _ANGLES = ['theta', 'phi', 'chi', 'zeta', 'gamma']
 
 
-def _create_tests(n, seed):
+def _create_tests(n, seed, with_options: bool = False):
     rng = np.random.default_rng(seed)
     angles = (rng.random((n, 5)) * 2 - 1) * np.pi
     # Add errors to the first 2 angles (theta and phi).
@@ -38,6 +38,23 @@ def _create_tests(n, seed):
     error = np.concatenate(
         [(rng.random((n, 2)) + 1) * rng.choice([-1, 1], (n, 2)), np.zeros((n, 3))], axis=-1
     )
+    if with_options:
+        options = []
+        for _ in range(n):
+            v = [False, False, False]
+            # Calibrate only one to keep the run time down.
+            v[rng.integers(0, 3)] = True
+            options.append(
+                {
+                    'characterize_chi': v[0],
+                    'characterize_gamma': v[1],
+                    'characterize_zeta': v[2],
+                    'characterize_phi': False,
+                    'characterize_theta': False,
+                }
+            )
+
+        return zip(angles, error, options)
     return zip(angles, error)
 
 
@@ -68,19 +85,17 @@ def _core_iterator(
         yield from super()._core_iterator(new_circuit, sim_state, all_measurements_are_terminal)
 
 
-@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=10, seed=32432432))
-def test_calibrate_z_phases(angles, error):
+@pytest.mark.parametrize(
+    ['angles', 'error', 'characterization_flags'],
+    _create_tests(n=10, seed=32432432, with_options=True),
+)
+def test_calibrate_z_phases(angles, error, characterization_flags):
 
     original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})
     actual_gate = cirq.PhasedFSimGate(**{k: v + e for k, v, e in zip(_ANGLES, angles, error)})
 
     options = XEBPhasedFSimCharacterizationOptions(
-        **{f'{n}_default': t for n, t in zip(_ANGLES, angles)},
-        characterize_chi=False,
-        characterize_gamma=False,
-        characterize_phi=True,
-        characterize_theta=True,
-        characterize_zeta=False,
+        **{f'{n}_default': t for n, t in zip(_ANGLES, angles)}, **characterization_flags
     )
 
     sampler = _TestSimulator(original_gate, actual_gate, seed=0)
@@ -111,7 +126,7 @@ def test_calibrate_z_phases(angles, error):
     assert new_dist < original_dist or new_dist < 1e-6
 
 
-@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=10, seed=32432432))
+@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=3, seed=32432432))
 def test_calibrate_z_phases_no_options(angles, error):
 
     original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})
@@ -145,7 +160,7 @@ def test_calibrate_z_phases_no_options(angles, error):
     assert new_dist < original_dist or new_dist < 1e-6
 
 
-@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=10, seed=32432432))
+@pytest.mark.parametrize(['angles', 'error'], _create_tests(n=3, seed=32432432))
 def test_calibrate_z_phases_workflow_no_options(angles, error):
 
     original_gate = cirq.PhasedFSimGate(**{k: v for k, v in zip(_ANGLES, angles)})
@@ -165,9 +180,11 @@ def test_calibrate_z_phases_workflow_no_options(angles, error):
     )
 
     for params in result.final_params.values():
-        assert 'zeta' not in params
-        assert 'chi' not in params
-        assert 'gamma' not in params
+        assert 'zeta' in params
+        assert 'chi' in params
+        assert 'gamma' in params
+        assert 'phi' not in params
+        assert 'theta' not in params
 
 
 def test_plot_z_phase_calibration_result():
@@ -180,8 +197,7 @@ def test_plot_z_phase_calibration_result():
     df['fidelities_c'] = [[0.9, 0.92, 0.93], [0.7, 0.77, 0.8]]
     df['layer_fid_std_c'] = [0.2, 0.3]
 
-    _, axes = plt.subplots(1, 2)
-    plot_z_phase_calibration_result(before_after_df=df, axes=axes)
+    axes = plot_z_phase_calibration_result(before_after_df=df)
 
     np.testing.assert_allclose(axes[0].lines[0].get_xdata().astype(float), [1, 2, 3])
     np.testing.assert_allclose(axes[0].lines[0].get_ydata().astype(float), [0.9, 0.8, 0.7])

From 55c47a5ea20e770f4deb8794740e41ec1065c02e Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Wed, 30 Oct 2024 18:49:06 -0700
Subject: [PATCH 16/22] address comments

---
 cirq-core/cirq/experiments/z_phase_calibration.py      | 7 ++++---
 cirq-core/cirq/experiments/z_phase_calibration_test.py | 7 +++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 6638689a162..f352adb5f9a 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/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 Callable, Union, Optional, Sequence, Tuple, Dict, TYPE_CHECKING
+from typing import Union, Optional, Sequence, Tuple, Dict, TYPE_CHECKING
 
 import multiprocessing
 import multiprocessing.pool
@@ -130,6 +130,7 @@ def z_phase_calibration_workflow(
     )
 
     if local_pool:
+        assert isinstance(pool, multiprocessing.pool.Pool)
         pool.close()
     return result, before_after
 
@@ -222,11 +223,11 @@ def calibrate_z_phases(
 
 def plot_z_phase_calibration_result(
     before_after_df: 'pd.DataFrame',
-    axes: Optional[Sequence['plt.Axes']] = None,
+    axes: Optional[np.ndarray[Sequence[Sequence['plt.Axes']], np.dtype[np.object_]]] = None,
     pairs: Optional[Sequence[Tuple['cirq.Qid', 'cirq.Qid']]] = None,
     *,
     with_error_bars: bool = False,
-) -> None:
+) -> np.ndarray[Sequence[Sequence['plt.Axes']], np.dtype[np.object_]]:
     """A helper method to plot the result of running z-phase calibration.
 
     Note that the plotted fidelity is a statistical estimate of the true fidelity and as a result
diff --git a/cirq-core/cirq/experiments/z_phase_calibration_test.py b/cirq-core/cirq/experiments/z_phase_calibration_test.py
index 2e29d48e8c8..c7c149b37c5 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration_test.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration_test.py
@@ -15,7 +15,6 @@
 import pytest
 import numpy as np
 import pandas as pd
-import matplotlib.pyplot as plt
 
 import cirq
 
@@ -32,11 +31,11 @@
 def _create_tests(n, seed, with_options: bool = False):
     rng = np.random.default_rng(seed)
     angles = (rng.random((n, 5)) * 2 - 1) * np.pi
-    # Add errors to the first 2 angles (theta and phi).
-    # The errors for theta and phi are in the union (-2, -1) U (1, 2).
+    # Add errors to the last 3 angles (chi, zeta, gamma).
+    # The errors are in the union (-2, -1) U (1, 2).
     # This is because we run the tests with few repetitions so a small error might not get fixed.
     error = np.concatenate(
-        [(rng.random((n, 2)) + 1) * rng.choice([-1, 1], (n, 2)), np.zeros((n, 3))], axis=-1
+        [np.zeros((n, 2)), (rng.random((n, 3)) + 1) * rng.choice([-1, 1], (n, 3))], axis=-1
     )
     if with_options:
         options = []

From 5d668e20e49d0cb770bbb348183483c07edaed74 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Thu, 7 Nov 2024 12:55:42 -0800
Subject: [PATCH 17/22] address comments

---
 cirq-core/cirq/experiments/__init__.py        |  5 ++-
 .../cirq/experiments/z_phase_calibration.py   | 32 ++++++++++---------
 2 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/cirq-core/cirq/experiments/__init__.py b/cirq-core/cirq/experiments/__init__.py
index 717ea658893..96310ca63f7 100644
--- a/cirq-core/cirq/experiments/__init__.py
+++ b/cirq-core/cirq/experiments/__init__.py
@@ -84,4 +84,7 @@
 )
 
 
-from cirq.experiments.z_phase_calibration import z_phase_calibration_workflow, calibrate_z_phases
+from cirq.experiments.z_phase_calibration import (
+    z_phase_calibration_workflow as z_phase_calibration_workflow,
+    calibrate_z_phases as calibrate_z_phases,
+)
diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index f352adb5f9a..97c9dabbede 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -14,9 +14,9 @@
 
 """Provides a method to do z-phase calibration for excitation-preserving gates."""
 from typing import Union, Optional, Sequence, Tuple, Dict, TYPE_CHECKING
-
 import multiprocessing
 import multiprocessing.pool
+
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -40,7 +40,7 @@ def z_phase_calibration_workflow(
     cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
-    pool_or_num_workers: Optional[Union[int, 'multiprocessing.pool.Pool']] = None,
+    num_workers_or_pool: Union[int, 'multiprocessing.pool.Pool'] = -1,
 ) -> Tuple[xeb_fitting.XEBCharacterizationResult, 'pd.DataFrame']:
     """Perform z-phase calibration for excitation-preserving gates.
 
@@ -64,8 +64,9 @@ def z_phase_calibration_workflow(
         sampler: The quantum engine or simulator to run the circuits.
         qubits: Qubits to use. If none, use all qubits on the sampler's device.
         two_qubit_gate: The entangling gate to use.
-        options: The XEB-fitting options. If None, calibrate all 5 PhasedFSimGate parameters,
-            using the representation of a two-qubit gate as an FSimGate for the initial guess.
+        options: The XEB-fitting options. If None, calibrate all only the three phase angles
+            (chi, gamma, zeta) using the representation of a two-qubit gate as an FSimGate
+            for the initial guess.
         n_repetitions: The number of repetitions to use.
         n_combinations: The number of combinations to generate.
         n_circuits: The number of circuits to generate.
@@ -75,7 +76,7 @@ def z_phase_calibration_workflow(
         pool_or_num_workers: An optional multi-processing pool or number of workers.
             A zero value means no multiprocessing.
             A positivie integer value will create a pool with the given number of workers.
-            A None value will create pool with maximum number of workers.
+            A negative value will create pool with maximum number of workers.
     Returns:
         - An `XEBCharacterizationResult` object that contains the calibration result.
         - A `pd.DataFrame` comparing the before and after fidelities.
@@ -83,10 +84,10 @@ def z_phase_calibration_workflow(
 
     pool: Optional['multiprocessing.pool.Pool'] = None
     local_pool = False
-    if isinstance(pool_or_num_workers, multiprocessing.pool.Pool):
-        pool = pool_or_num_workers  # pragma: no cover
-    elif pool_or_num_workers != 0:
-        pool = multiprocessing.Pool(pool_or_num_workers)
+    if isinstance(num_workers_or_pool, multiprocessing.pool.Pool):
+        pool = num_workers_or_pool  # pragma: no cover
+    elif num_workers_or_pool != 0:
+        pool = multiprocessing.Pool(num_workers_or_pool if num_workers_or_pool > 0 else None)
         local_pool = True
 
     fids_df_0, circuits, sampled_df = parallel_xeb_workflow(
@@ -146,7 +147,7 @@ def calibrate_z_phases(
     cycle_depths: Sequence[int] = tuple(np.arange(3, 100, 20)),
     random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     atol: float = 1e-3,
-    pool_or_num_workers: Optional[Union[int, 'multiprocessing.pool.Pool']] = None,
+    num_workers_or_pool: Union[int, 'multiprocessing.pool.Pool'] = -1,
 ) -> Dict[Tuple['cirq.Qid', 'cirq.Qid'], 'cirq.PhasedFSimGate']:
     """Perform z-phase calibration for excitation-preserving gates.
 
@@ -170,18 +171,19 @@ def calibrate_z_phases(
         sampler: The quantum engine or simulator to run the circuits.
         qubits: Qubits to use. If none, use all qubits on the sampler's device.
         two_qubit_gate: The entangling gate to use.
-        options: The XEB-fitting options. If None, calibrate all 5 PhasedFSimGate parameters,
-            using the representation of a two-qubit gate as an FSimGate for the initial guess.
+        options: The XEB-fitting options. If None, calibrate all only the three phase angles
+            (chi, gamma, zeta) using the representation of a two-qubit gate as an FSimGate
+            for the initial guess.
         n_repetitions: The number of repetitions to use.
         n_combinations: The number of combinations to generate.
         n_circuits: The number of circuits to generate.
         cycle_depths: The cycle depths to use.
         random_state: The random state to use.
         atol: Absolute tolerance to be used by the minimizer.
-        pool_or_num_workers: An optional multi-processing pool or number of workers.
+        num_workers_or_pool: An optional multi-processing pool or number of workers.
             A zero value means no multiprocessing.
             A positivie integer value will create a pool with the given number of workers.
-            A None value will create pool with maximum number of workers.
+            A negative value will create pool with maximum number of workers.
 
     Returns:
         - A dictionary mapping qubit pairs to the calibrated PhasedFSimGates.
@@ -207,7 +209,7 @@ def calibrate_z_phases(
         cycle_depths=cycle_depths,
         random_state=random_state,
         atol=atol,
-        pool_or_num_workers=pool_or_num_workers,
+        pool_or_num_workers=num_workers_or_pool,
     )
 
     gates = {}

From a372c8e8755127e51f85a6ed38c4262ef8d338c1 Mon Sep 17 00:00:00 2001
From: Nour Yosri <noureldinyosri@gmail.com>
Date: Thu, 7 Nov 2024 13:44:29 -0800
Subject: [PATCH 18/22] nit

---
 cirq-core/cirq/experiments/z_phase_calibration.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 97c9dabbede..0c1de94f30a 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -73,7 +73,7 @@ def z_phase_calibration_workflow(
         cycle_depths: The cycle depths to use.
         random_state: The random state to use.
         atol: Absolute tolerance to be used by the minimizer.
-        pool_or_num_workers: An optional multi-processing pool or number of workers.
+        num_workers_or_pool: An optional multi-processing pool or number of workers.
             A zero value means no multiprocessing.
             A positivie integer value will create a pool with the given number of workers.
             A negative value will create pool with maximum number of workers.
@@ -209,7 +209,7 @@ def calibrate_z_phases(
         cycle_depths=cycle_depths,
         random_state=random_state,
         atol=atol,
-        pool_or_num_workers=num_workers_or_pool,
+        num_workers_or_pool=num_workers_or_pool,
     )
 
     gates = {}

From d9a82828e18597e7f042e9369c902919fe5a0fca Mon Sep 17 00:00:00 2001
From: Noureldin <noureldinyosri@gmail.com>
Date: Thu, 7 Nov 2024 15:55:01 -0800
Subject: [PATCH 19/22] Update
 cirq-core/cirq/experiments/z_phase_calibration.py

Co-authored-by: Pavol Juhas <pavol.juhas@gmail.com>
---
 cirq-core/cirq/experiments/z_phase_calibration.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 0c1de94f30a..ed8059af2f8 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -182,7 +182,7 @@ def calibrate_z_phases(
         atol: Absolute tolerance to be used by the minimizer.
         num_workers_or_pool: An optional multi-processing pool or number of workers.
             A zero value means no multiprocessing.
-            A positivie integer value will create a pool with the given number of workers.
+            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.
 
     Returns:

From b19cee9d614e0db6fe69d8d9a193fdc0a056deb7 Mon Sep 17 00:00:00 2001
From: Noureldin <noureldinyosri@gmail.com>
Date: Thu, 7 Nov 2024 15:55:09 -0800
Subject: [PATCH 20/22] Update
 cirq-core/cirq/experiments/z_phase_calibration.py

Co-authored-by: Pavol Juhas <pavol.juhas@gmail.com>
---
 cirq-core/cirq/experiments/z_phase_calibration.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index ed8059af2f8..33151440249 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -171,7 +171,7 @@ def calibrate_z_phases(
         sampler: The quantum engine or simulator to run the circuits.
         qubits: Qubits to use. If none, use all qubits on the sampler's device.
         two_qubit_gate: The entangling gate to use.
-        options: The XEB-fitting options. If None, calibrate all only the three phase angles
+        options: The XEB-fitting options. If None, calibrate only the three phase angles
             (chi, gamma, zeta) using the representation of a two-qubit gate as an FSimGate
             for the initial guess.
         n_repetitions: The number of repetitions to use.

From 360aee17ac19d80ed9c41e2f79d14e1859a1e127 Mon Sep 17 00:00:00 2001
From: Noureldin <noureldinyosri@gmail.com>
Date: Thu, 7 Nov 2024 15:55:17 -0800
Subject: [PATCH 21/22] Update
 cirq-core/cirq/experiments/z_phase_calibration.py

Co-authored-by: Pavol Juhas <pavol.juhas@gmail.com>
---
 cirq-core/cirq/experiments/z_phase_calibration.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index 33151440249..e23ec817f31 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -75,7 +75,7 @@ def z_phase_calibration_workflow(
         atol: Absolute tolerance to be used by the minimizer.
         num_workers_or_pool: An optional multi-processing pool or number of workers.
             A zero value means no multiprocessing.
-            A positivie integer value will create a pool with the given number of workers.
+            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.
     Returns:
         - An `XEBCharacterizationResult` object that contains the calibration result.

From 296922d80c03215554264adf223a3a1499b62291 Mon Sep 17 00:00:00 2001
From: Noureldin <noureldinyosri@gmail.com>
Date: Thu, 7 Nov 2024 15:55:25 -0800
Subject: [PATCH 22/22] Update
 cirq-core/cirq/experiments/z_phase_calibration.py

Co-authored-by: Pavol Juhas <pavol.juhas@gmail.com>
---
 cirq-core/cirq/experiments/z_phase_calibration.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cirq-core/cirq/experiments/z_phase_calibration.py b/cirq-core/cirq/experiments/z_phase_calibration.py
index e23ec817f31..363810fbd13 100644
--- a/cirq-core/cirq/experiments/z_phase_calibration.py
+++ b/cirq-core/cirq/experiments/z_phase_calibration.py
@@ -64,7 +64,7 @@ def z_phase_calibration_workflow(
         sampler: The quantum engine or simulator to run the circuits.
         qubits: Qubits to use. If none, use all qubits on the sampler's device.
         two_qubit_gate: The entangling gate to use.
-        options: The XEB-fitting options. If None, calibrate all only the three phase angles
+        options: The XEB-fitting options. If None, calibrate only the three phase angles
             (chi, gamma, zeta) using the representation of a two-qubit gate as an FSimGate
             for the initial guess.
         n_repetitions: The number of repetitions to use.