cirq-core 1.5.0.dev20250409225226__py3-none-any.whl → 1.6.0__py3-none-any.whl

cirq/work/observable_measurement.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import abc
 import dataclasses
 import itertools
 import os
 import tempfile
 import warnings
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
+from typing import Any, Iterable, Mapping, Sequence, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -50,8 +52,8 @@ document(
 
 
 def _with_parameterized_layers(
-    circuit: 'cirq.AbstractCircuit', qubits: Sequence['cirq.Qid'], needs_init_layer: bool
-) -> 'cirq.Circuit':
+    circuit: cirq.AbstractCircuit, qubits: Sequence[cirq.Qid], needs_init_layer: bool
+) -> cirq.Circuit:
     """Return a copy of the input circuit with parameterized single-qubit rotations.
 
     These rotations flank the circuit: the initial two layers of X and Y gates
@@ -136,7 +138,7 @@ class RepetitionsStoppingCriteria(StoppingCriteria):
         return protocols.dataclass_json_dict(self)
 
 
-_OBS_TO_PARAM_VAL: Dict[Tuple['cirq.Pauli', bool], Tuple[float, float]] = {
+_OBS_TO_PARAM_VAL: dict[tuple[cirq.Pauli, bool], tuple[float, float]] = {
     (ops.X, False): (0, -1 / 2),
     (ops.X, True): (0, +1 / 2),
     (ops.Y, False): (1 / 2, 0),
@@ -148,7 +150,7 @@ _OBS_TO_PARAM_VAL: Dict[Tuple['cirq.Pauli', bool], Tuple[float, float]] = {
 second element in the key is whether to measure in the positive or negative (flipped) basis
 for readout symmetrization."""
 
-_STATE_TO_PARAM_VAL: Dict['_NamedOneQubitState', Tuple[float, float]] = {
+_STATE_TO_PARAM_VAL: dict[_NamedOneQubitState, tuple[float, float]] = {
     value.KET_PLUS: (0, +1 / 2),
     value.KET_MINUS: (0, -1 / 2),
     value.KET_IMAG: (-1 / 2, 0),
@@ -162,9 +164,9 @@ _STATE_TO_PARAM_VAL: Dict['_NamedOneQubitState', Tuple[float, float]] = {
 def _get_params_for_setting(
     setting: InitObsSetting,
     flips: Iterable[bool],
-    qubits: Sequence['cirq.Qid'],
+    qubits: Sequence[cirq.Qid],
     needs_init_layer: bool,
-) -> Dict[str, float]:
+) -> dict[str, float]:
     """Return the parameter dictionary for the given setting.
 
     This must be used in conjunction with a circuit generated by
@@ -204,9 +206,9 @@ def _get_params_for_setting(
 
 def _pad_setting(
     max_setting: InitObsSetting,
-    qubits: Sequence['cirq.Qid'],
+    qubits: Sequence[cirq.Qid],
     pad_init_state_with=value.KET_ZERO,
-    pad_obs_with: 'cirq.Gate' = ops.Z,
+    pad_obs_with: cirq.Gate = ops.Z,
 ) -> InitObsSetting:
     """Pad `max_setting`'s `init_state` and `observable` with `pad_xx_with` operations
     (defaults:  |0> and Z) so each max_setting has the same qubits. We need this
@@ -227,7 +229,7 @@ def _pad_setting(
     return InitObsSetting(init_state=init_state, observable=obs)
 
 
-def _aggregate_n_repetitions(next_chunk_repetitions: Set[int]) -> int:
+def _aggregate_n_repetitions(next_chunk_repetitions: set[int]) -> int:
     """A stopping criteria can request a different number of more_repetitions for each
     measurement spec. For batching efficiency, we take the max and issue a warning in this case."""
     if len(next_chunk_repetitions) == 1:
@@ -243,10 +245,10 @@ def _aggregate_n_repetitions(next_chunk_repetitions: Set[int]) -> int:
 
 
 def _check_meas_specs_still_todo(
-    meas_specs: List[_MeasurementSpec],
-    accumulators: Dict[_MeasurementSpec, BitstringAccumulator],
+    meas_specs: list[_MeasurementSpec],
+    accumulators: Mapping[_MeasurementSpec, BitstringAccumulator],
     stopping_criteria: StoppingCriteria,
-) -> Tuple[List[_MeasurementSpec], int]:
+) -> tuple[list[_MeasurementSpec], int]:
     """Filter `meas_specs` in case some are done.
 
     In the sampling loop in `measure_grouped_settings`, we submit
@@ -254,7 +256,7 @@ def _check_meas_specs_still_todo(
     removing `meas_spec`s from the loop if they are done.
     """
     still_todo = []
-    repetitions_set: Set[int] = set()
+    repetitions_set: set[int] = set()
     for meas_spec in meas_specs:
         accumulator = accumulators[meas_spec]
         more_repetitions = stopping_criteria.more_repetitions(accumulator)
@@ -303,7 +305,7 @@ class _FlippyMeasSpec:
 
     meas_spec: _MeasurementSpec
     flips: np.ndarray
-    qubits: Sequence['cirq.Qid']
+    qubits: Sequence[cirq.Qid]
 
     def param_tuples(self, *, needs_init_layer=True):
         yield from _get_params_for_setting(
@@ -318,9 +320,9 @@ class _FlippyMeasSpec:
 def _subdivide_meas_specs(
     meas_specs: Iterable[_MeasurementSpec],
     repetitions: int,
-    qubits: Sequence['cirq.Qid'],
+    qubits: Sequence[cirq.Qid],
     readout_symmetrization: bool,
-) -> Tuple[List[_FlippyMeasSpec], int]:
+) -> tuple[list[_FlippyMeasSpec], int]:
     """Split measurement specs into sub-jobs for readout symmetrization
 
     In readout symmetrization, we first run the "normal" circuit followed
@@ -354,8 +356,8 @@ def _to_sweep(param_tuples):
 
 
 def _parse_checkpoint_options(
-    checkpoint: bool, checkpoint_fn: Optional[str], checkpoint_other_fn: Optional[str]
-) -> Tuple[Optional[str], Optional[str]]:
+    checkpoint: bool, checkpoint_fn: str | None, checkpoint_other_fn: str | None
+) -> tuple[str | None, str | None]:
     """Parse the checkpoint-oriented options in `measure_grouped_settings`.
 
     This function contains the validation and defaults logic. Please see
@@ -433,8 +435,8 @@ class CheckpointFileOptions:
     """
 
     checkpoint: bool = False
-    checkpoint_fn: Optional[str] = None
-    checkpoint_other_fn: Optional[str] = None
+    checkpoint_fn: str | None = None
+    checkpoint_other_fn: str | None = None
 
     def __post_init__(self):
         fn, other_fn = _parse_checkpoint_options(
@@ -458,7 +460,7 @@ class CheckpointFileOptions:
         protocols.to_json(obj, self.checkpoint_fn)
 
 
-def _needs_init_layer(grouped_settings: Dict[InitObsSetting, List[InitObsSetting]]) -> bool:
+def _needs_init_layer(grouped_settings: dict[InitObsSetting, list[InitObsSetting]]) -> bool:
     """Helper function to go through init_states and determine if any of them need an
     initialization layer of single-qubit gates."""
     for max_setting in grouped_settings.keys():
@@ -468,16 +470,16 @@ def _needs_init_layer(grouped_settings: Dict[InitObsSetting, List[InitObsSetting
 
 
 def measure_grouped_settings(
-    circuit: 'cirq.AbstractCircuit',
-    grouped_settings: Dict[InitObsSetting, List[InitObsSetting]],
-    sampler: 'cirq.Sampler',
+    circuit: cirq.AbstractCircuit,
+    grouped_settings: dict[InitObsSetting, list[InitObsSetting]],
+    sampler: cirq.Sampler,
     stopping_criteria: StoppingCriteria,
     *,
     readout_symmetrization: bool = False,
-    circuit_sweep: 'cirq.Sweepable' = None,
-    readout_calibrations: Optional[BitstringAccumulator] = None,
+    circuit_sweep: cirq.Sweepable = None,
+    readout_calibrations: BitstringAccumulator | None = None,
     checkpoint: CheckpointFileOptions = CheckpointFileOptions(),
-) -> List[BitstringAccumulator]:
+) -> list[BitstringAccumulator]:
     """Measure a suite of grouped InitObsSetting settings.
 
     This is a low-level API for accessing the observable measurement
@@ -582,10 +584,10 @@ def measure_grouped_settings(
     return list(accumulators.values())
 
 
-_GROUPING_FUNCS: Dict[str, GROUPER_T] = {'greedy': group_settings_greedy}
+_GROUPING_FUNCS: dict[str, GROUPER_T] = {'greedy': group_settings_greedy}
 
 
-def _parse_grouper(grouper: Union[str, GROUPER_T] = group_settings_greedy) -> GROUPER_T:
+def _parse_grouper(grouper: str | GROUPER_T = group_settings_greedy) -> GROUPER_T:
     """Logic for turning a named grouper into one of the build-in groupers in support of the
     high-level `measure_observables` API."""
     if isinstance(grouper, str):
@@ -597,8 +599,8 @@ def _parse_grouper(grouper: Union[str, GROUPER_T] = group_settings_greedy) -> GR
 
 
 def _get_all_qubits(
-    circuit: 'cirq.AbstractCircuit', observables: Iterable['cirq.PauliString']
-) -> List['cirq.Qid']:
+    circuit: cirq.AbstractCircuit, observables: Iterable[cirq.PauliString]
+) -> list[cirq.Qid]:
     """Helper function for `measure_observables` to get all qubits from a circuit and a
     collection of observables."""
     qubit_set = set()
@@ -609,17 +611,17 @@ def _get_all_qubits(
 
 
 def measure_observables(
-    circuit: 'cirq.AbstractCircuit',
-    observables: Iterable['cirq.PauliString'],
-    sampler: Union['cirq.Simulator', 'cirq.Sampler'],
+    circuit: cirq.AbstractCircuit,
+    observables: Iterable[cirq.PauliString],
+    sampler: cirq.Simulator | cirq.Sampler,
     stopping_criteria: StoppingCriteria,
     *,
     readout_symmetrization: bool = False,
-    circuit_sweep: Optional['cirq.Sweepable'] = None,
-    grouper: Union[str, GROUPER_T] = group_settings_greedy,
-    readout_calibrations: Optional[BitstringAccumulator] = None,
+    circuit_sweep: cirq.Sweepable | None = None,
+    grouper: str | GROUPER_T = group_settings_greedy,
+    readout_calibrations: BitstringAccumulator | None = None,
     checkpoint: CheckpointFileOptions = CheckpointFileOptions(),
-) -> List[ObservableMeasuredResult]:
+) -> list[ObservableMeasuredResult]:
     """Measure a collection of PauliString observables for a state prepared by a Circuit.
 
     If you need more control over the process, please see `measure_grouped_settings` for a
@@ -671,15 +673,15 @@ def measure_observables(
 
 
 def measure_observables_df(
-    circuit: 'cirq.AbstractCircuit',
-    observables: Iterable['cirq.PauliString'],
-    sampler: Union['cirq.Simulator', 'cirq.Sampler'],
+    circuit: cirq.AbstractCircuit,
+    observables: Iterable[cirq.PauliString],
+    sampler: cirq.Simulator | cirq.Sampler,
     stopping_criteria: StoppingCriteria,
     *,
     readout_symmetrization: bool = False,
-    circuit_sweep: Optional['cirq.Sweepable'] = None,
-    grouper: Union[str, GROUPER_T] = group_settings_greedy,
-    readout_calibrations: Optional[BitstringAccumulator] = None,
+    circuit_sweep: cirq.Sweepable | None = None,
+    grouper: str | GROUPER_T = group_settings_greedy,
+    readout_calibrations: BitstringAccumulator | None = None,
     checkpoint: CheckpointFileOptions = CheckpointFileOptions(),
 ):
     """Measure observables and return resulting data as a Pandas dataframe.

cirq/work/observable_measurement_data.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import dataclasses
 import datetime
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, TYPE_CHECKING, Union
+from typing import Any, Iterable, Mapping, TYPE_CHECKING
 
 import numpy as np
 import sympy
@@ -33,7 +35,7 @@ if TYPE_CHECKING:
     import cirq
 
 
-def _check_and_get_real_coef(observable: 'cirq.PauliString', atol: float):
+def _check_and_get_real_coef(observable: cirq.PauliString, atol: float):
     """Assert that a PauliString has a real coefficient and return it."""
     coef = observable.coefficient
     if isinstance(coef, sympy.Expr) or not np.isclose(coef.imag, 0, atol=atol):
@@ -43,8 +45,8 @@ def _check_and_get_real_coef(observable: 'cirq.PauliString', atol: float):
 
 def _obs_vals_from_measurements(
     bitstrings: np.ndarray,
-    qubit_to_index: Dict['cirq.Qid', int],
-    observable: 'cirq.PauliString',
+    qubit_to_index: Mapping[cirq.Qid, int],
+    observable: cirq.PauliString,
     atol: float,
 ):
     """Multiply together bitstrings to get observed values of operators."""
@@ -61,10 +63,10 @@ def _obs_vals_from_measurements(
 
 def _stats_from_measurements(
     bitstrings: np.ndarray,
-    qubit_to_index: Dict['cirq.Qid', int],
-    observable: 'cirq.PauliString',
+    qubit_to_index: Mapping[cirq.Qid, int],
+    observable: cirq.PauliString,
     atol: float,
-) -> Tuple[float, float]:
+) -> tuple[float, float]:
     """Return the mean and squared standard error of the mean for the given
     observable according to the measurements in `bitstrings`."""
     obs_vals = _obs_vals_from_measurements(bitstrings, qubit_to_index, observable, atol=atol)
@@ -107,7 +109,7 @@ class ObservableMeasuredResult:
     mean: float
     variance: float
     repetitions: int
-    circuit_params: Mapping[Union[str, sympy.Expr], Union[value.Scalar, sympy.Expr]]
+    circuit_params: Mapping[str | sympy.Expr, value.Scalar | sympy.Expr]
 
     # unhashable because of the mapping-type circuit_params attribute
     __hash__ = None  # type: ignore
@@ -136,7 +138,7 @@ class ObservableMeasuredResult:
     def stddev(self):
         return np.sqrt(self.variance)
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Return the contents of this class as a dictionary.
 
         This makes records suitable for construction of a Pandas dataframe. The circuit parameters
@@ -213,29 +215,32 @@ class BitstringAccumulator:
     def __init__(
         self,
         meas_spec: _MeasurementSpec,
-        simul_settings: List[InitObsSetting],
-        qubit_to_index: Dict['cirq.Qid', int],
-        bitstrings: Optional[np.ndarray] = None,
-        chunksizes: Optional[np.ndarray] = None,
-        timestamps: Optional[np.ndarray] = None,
-        readout_calibration: Optional['BitstringAccumulator'] = None,
+        simul_settings: list[InitObsSetting],
+        qubit_to_index: Mapping[cirq.Qid, int],
+        bitstrings: np.ndarray | None = None,
+        chunksizes: np.ndarray | None = None,
+        timestamps: np.ndarray | None = None,
+        readout_calibration: BitstringAccumulator | None = None,
     ):
         self._meas_spec = meas_spec
         self._simul_settings = simul_settings
         self._qubit_to_index = qubit_to_index
         self._readout_calibration = readout_calibration
 
+        self.bitstrings: np.ndarray[tuple[int, ...], np.dtype[np.uint8]]
         if bitstrings is None:
             n_bits = len(qubit_to_index)
             self.bitstrings = np.zeros((0, n_bits), dtype=np.uint8)
         else:
             self.bitstrings = np.asarray(bitstrings, dtype=np.uint8)
 
+        self.chunksizes: np.ndarray[tuple[int, ...], np.dtype[np.int64]]
         if chunksizes is None:
             self.chunksizes = np.zeros((0,), dtype=np.int64)
         else:
             self.chunksizes = np.asarray(chunksizes, dtype=np.int64)
 
+        self.timestamps: np.ndarray[tuple[int, ...], np.dtype[np.datetime64]]
         if timestamps is None:
             self.timestamps = np.zeros((0,), dtype='datetime64[us]')
         else:
@@ -523,8 +528,8 @@ class BitstringAccumulator:
 
 
 def flatten_grouped_results(
-    grouped_results: List[BitstringAccumulator],
-) -> List[ObservableMeasuredResult]:
+    grouped_results: list[BitstringAccumulator],
+) -> list[ObservableMeasuredResult]:
     """Flatten a collection of BitstringAccumulators into a list of ObservableMeasuredResult.
 
     Raw results are contained in BitstringAccumulator which contains

cirq/work/observable_measurement_data_test.py

@@ -11,6 +11,9 @@
 # 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.
+
+from __future__ import annotations
+
 import dataclasses
 import datetime
 import time
@@ -113,7 +116,7 @@ def test_observable_measured_result():
 
 
 @pytest.fixture
-def example_bsa() -> 'cw.BitstringAccumulator':
+def example_bsa() -> cw.BitstringAccumulator:
     """Test fixture to create an (empty) example BitstringAccumulator"""
     q0, q1 = cirq.LineQubit.range(2)
     setting = cw.InitObsSetting(

cirq/work/observable_measurement_test.py

@@ -11,11 +11,15 @@
 # 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.
+
+from __future__ import annotations
+
 import tempfile
-from typing import Dict, Iterable, List
+from typing import Iterable
 
 import numpy as np
 import pytest
+import sympy
 
 import cirq
 import cirq.work as cw
@@ -37,7 +41,7 @@ from cirq.work.observable_measurement import (
 )
 
 
-def test_with_parameterized_layers():
+def test_with_parameterized_layers() -> None:
     qs = cirq.LineQubit.range(3)
     circuit = cirq.Circuit([cirq.H.on_each(*qs), cirq.CZ(qs[0], qs[1]), cirq.CZ(qs[1], qs[2])])
     circuit2 = _with_parameterized_layers(circuit, qubits=qs, needs_init_layer=False)
@@ -46,9 +50,11 @@ def test_with_parameterized_layers():
     *_, xlayer, ylayer, measurelayer = circuit2.moments
     for op in xlayer.operations:
         assert isinstance(op.gate, cirq.XPowGate)
+        assert isinstance(op.gate.exponent, sympy.Symbol)
         assert op.gate.exponent.name.endswith('-Xf')
     for op in ylayer.operations:
         assert isinstance(op.gate, cirq.YPowGate)
+        assert isinstance(op.gate.exponent, sympy.Symbol)
         assert op.gate.exponent.name.endswith('-Yf')
     for op in measurelayer:
         assert isinstance(op.gate, cirq.MeasurementGate)
@@ -60,13 +66,15 @@ def test_with_parameterized_layers():
     xlayer, ylayer, *_ = circuit3.moments
     for op in xlayer.operations:
         assert isinstance(op.gate, cirq.XPowGate)
+        assert isinstance(op.gate.exponent, sympy.Symbol)
         assert op.gate.exponent.name.endswith('-Xi')
     for op in ylayer.operations:
         assert isinstance(op.gate, cirq.YPowGate)
+        assert isinstance(op.gate.exponent, sympy.Symbol)
         assert op.gate.exponent.name.endswith('-Yi')
 
 
-def test_get_params_for_setting():
+def test_get_params_for_setting() -> None:
     qubits = cirq.LineQubit.range(3)
     a, b, c = qubits
 
@@ -82,10 +90,10 @@ def test_get_params_for_setting():
     needs_init_layer = True
     with pytest.raises(ValueError):
         _get_params_for_setting(
-            padded_setting, flips=[0, 0], qubits=qubits, needs_init_layer=needs_init_layer
+            padded_setting, flips=[False, False], qubits=qubits, needs_init_layer=needs_init_layer
         )
     params = _get_params_for_setting(
-        padded_setting, flips=[0, 0, 1], qubits=qubits, needs_init_layer=needs_init_layer
+        padded_setting, flips=[False, False, True], qubits=qubits, needs_init_layer=needs_init_layer
     )
     assert all(
         x in params
@@ -116,7 +124,7 @@ def test_get_params_for_setting():
     np.testing.assert_allclose([ma, mb, mc], [1, 0, -1])
 
 
-def test_params_and_settings():
+def test_params_and_settings() -> None:
     qubits = cirq.LineQubit.range(1)
     (q,) = qubits
     tests = [
@@ -143,7 +151,7 @@ def test_params_and_settings():
         assert np.abs(coef - z) < 1e-2, f'{init} {obs} {coef}'
 
 
-def test_subdivide_meas_specs():
+def test_subdivide_meas_specs() -> None:
     qubits = cirq.LineQubit.range(2)
     q0, q1 = qubits
     setting = cw.InitObsSetting(
@@ -177,7 +185,7 @@ def test_subdivide_meas_specs():
     ]
 
 
-def test_aggregate_n_repetitions():
+def test_aggregate_n_repetitions() -> None:
     with pytest.warns(UserWarning):
         reps = _aggregate_n_repetitions({5, 6})
     assert reps == 6
@@ -197,7 +205,7 @@ class _MockBitstringAccumulator(BitstringAccumulator):
         return cov / len(self.bitstrings)
 
 
-def test_variance_stopping_criteria():
+def test_variance_stopping_criteria() -> None:
     stop = cw.VarianceStoppingCriteria(variance_bound=1e-6)
     acc = _MockBitstringAccumulator()
     assert stop.more_repetitions(acc) == 10_000
@@ -222,22 +230,30 @@ class _WildVarianceStoppingCriteria(StoppingCriteria):
         return [5, 6][self._state % 2]
 
 
-def test_variance_stopping_criteria_aggregate_n_repetitions():
+def test_variance_stopping_criteria_aggregate_n_repetitions() -> None:
+    q0, q1 = cirq.LineQubit.range(2)
     stop = _WildVarianceStoppingCriteria()
     acc1 = _MockBitstringAccumulator()
     acc2 = _MockBitstringAccumulator()
-    accumulators = {'FakeMeasSpec1': acc1, 'FakeMeasSpec2': acc2}
+    setting = InitObsSetting(
+        init_state=cirq.KET_ZERO(q0) * cirq.KET_ZERO(q1), observable=cirq.X(q0) * cirq.Y(q1)
+    )
+    meas_spec = _MeasurementSpec(
+        max_setting=setting, circuit_params={'beta': 0.123, 'gamma': 0.456}
+    )
+    meas_spec2 = _MeasurementSpec(
+        max_setting=setting, circuit_params={'beta': 0.123, 'gamma': 0.456}
+    )
+    accumulators = {meas_spec: acc1, meas_spec2: acc2}
     with pytest.warns(UserWarning, match='the largest value will be used: 6.'):
         still_todo, reps = _check_meas_specs_still_todo(
-            meas_specs=sorted(accumulators.keys()),
-            accumulators=accumulators,
-            stopping_criteria=stop,
+            meas_specs=[meas_spec, meas_spec2], accumulators=accumulators, stopping_criteria=stop
         )
-    assert still_todo == ['FakeMeasSpec1', 'FakeMeasSpec2']
+    assert still_todo == [meas_spec, meas_spec2]
     assert reps == 6
 
 
-def test_repetitions_stopping_criteria():
+def test_repetitions_stopping_criteria() -> None:
     stop = cw.RepetitionsStoppingCriteria(total_repetitions=50_000)
     acc = _MockBitstringAccumulator()
 
@@ -248,7 +264,7 @@ def test_repetitions_stopping_criteria():
     assert todos == [10_000] * 5 + [0, 0]
 
 
-def test_repetitions_stopping_criteria_partial():
+def test_repetitions_stopping_criteria_partial() -> None:
     stop = cw.RepetitionsStoppingCriteria(total_repetitions=5_000, repetitions_per_chunk=1_000_000)
     acc = _MockBitstringAccumulator()
     assert stop.more_repetitions(acc) == 5_000
@@ -268,7 +284,7 @@ def _set_up_meas_specs_for_testing():
     return bsa, meas_spec
 
 
-def test_meas_specs_still_todo():
+def test_meas_specs_still_todo() -> None:
     bsa, meas_spec = _set_up_meas_specs_for_testing()
     stop = cw.RepetitionsStoppingCriteria(1_000)
 
@@ -296,7 +312,7 @@ def test_meas_specs_still_todo():
     assert reps == 0
 
 
-def test_meas_spec_still_todo_bad_spec():
+def test_meas_spec_still_todo_bad_spec() -> None:
     bsa, meas_spec = _set_up_meas_specs_for_testing()
 
     class BadStopping(StoppingCriteria):
@@ -310,7 +326,7 @@ def test_meas_spec_still_todo_bad_spec():
         )
 
 
-def test_meas_spec_still_todo_too_many_params(monkeypatch):
+def test_meas_spec_still_todo_too_many_params(monkeypatch) -> None:
     monkeypatch.setattr(cw.observable_measurement, 'MAX_REPETITIONS_PER_JOB', 30_000)
     bsa, meas_spec = _set_up_meas_specs_for_testing()
     lots_of_meas_spec = [meas_spec] * 3_001
@@ -321,7 +337,7 @@ def test_meas_spec_still_todo_too_many_params(monkeypatch):
         )
 
 
-def test_meas_spec_still_todo_lots_of_params(monkeypatch):
+def test_meas_spec_still_todo_lots_of_params(monkeypatch) -> None:
     monkeypatch.setattr(cw.observable_measurement, 'MAX_REPETITIONS_PER_JOB', 30_000)
     bsa, meas_spec = _set_up_meas_specs_for_testing()
     lots_of_meas_spec = [meas_spec] * 4
@@ -332,7 +348,7 @@ def test_meas_spec_still_todo_lots_of_params(monkeypatch):
         )
 
 
-def test_checkpoint_options():
+def test_checkpoint_options() -> None:
     # There are three ~binary options (the latter two can be either specified or `None`. We
     # test those 2^3 cases.
 
@@ -345,12 +361,15 @@ def test_checkpoint_options():
         _parse_checkpoint_options(False, 'test1', 'test2')
 
     chk, chkprev = _parse_checkpoint_options(True, None, None)
+    assert chk is not None
+    assert chkprev is not None
     assert chk.startswith(tempfile.gettempdir())
     assert chk.endswith('observables.json')
     assert chkprev.startswith(tempfile.gettempdir())
     assert chkprev.endswith('observables.prev.json')
 
     chk, chkprev = _parse_checkpoint_options(True, None, 'prev.json')
+    assert chk is not None
     assert chk.startswith(tempfile.gettempdir())
     assert chk.endswith('observables.json')
     assert chkprev == 'prev.json'
@@ -381,7 +400,7 @@ def test_checkpoint_options():
 
 
 @pytest.mark.parametrize(('with_circuit_sweep', 'checkpoint'), [(True, True), (False, False)])
-def test_measure_grouped_settings(with_circuit_sweep, checkpoint, tmpdir):
+def test_measure_grouped_settings(with_circuit_sweep, checkpoint, tmpdir) -> None:
     qubits = cirq.LineQubit.range(1)
     (q,) = qubits
     tests = [
@@ -431,7 +450,7 @@ def _get_some_grouped_settings():
     return grouped_settings, qubits
 
 
-def test_measure_grouped_settings_calibration_validation():
+def test_measure_grouped_settings_calibration_validation() -> None:
     mock_ro_calib = _MockBitstringAccumulator()
     grouped_settings, qubits = _get_some_grouped_settings()
 
@@ -448,7 +467,7 @@ def test_measure_grouped_settings_calibration_validation():
         )
 
 
-def test_measure_grouped_settings_read_checkpoint(tmpdir):
+def test_measure_grouped_settings_read_checkpoint(tmpdir) -> None:
     qubits = cirq.LineQubit.range(1)
     (q,) = qubits
 
@@ -495,7 +514,7 @@ Q = cirq.NamedQubit('q')
         (cirq.Circuit(cirq.Y(Q) ** 0.5, cirq.Z(Q) ** 0.2), cirq.X(Q)),
     ],
 )
-def test_XYZ_point8(circuit, observable):
+def test_XYZ_point8(circuit, observable) -> None:
     # each circuit, observable combination should result in the observable value of 0.8
     df = measure_observables_df(
         circuit,
@@ -510,14 +529,14 @@ def test_XYZ_point8(circuit, observable):
 
 def _each_in_its_own_group_grouper(
     settings: Iterable[InitObsSetting],
-) -> Dict[InitObsSetting, List[InitObsSetting]]:
+) -> dict[InitObsSetting, list[InitObsSetting]]:
     return {setting: [setting] for setting in settings}
 
 
 @pytest.mark.parametrize(
     'grouper', ['greedy', group_settings_greedy, _each_in_its_own_group_grouper]
 )
-def test_measure_observable_grouper(grouper):
+def test_measure_observable_grouper(grouper) -> None:
     circuit = cirq.Circuit(cirq.X(Q) ** 0.2)
     observables = [cirq.Z(Q), cirq.Z(cirq.NamedQubit('q2'))]
     results = measure_observables(
@@ -532,7 +551,7 @@ def test_measure_observable_grouper(grouper):
     np.testing.assert_allclose(1, results[1].mean, atol=1e-9)
 
 
-def test_measure_observable_bad_grouper():
+def test_measure_observable_bad_grouper() -> None:
     circuit = cirq.Circuit(cirq.X(Q) ** 0.2)
     observables = [cirq.Z(Q), cirq.Z(cirq.NamedQubit('q2'))]
     with pytest.raises(ValueError, match=r'Unknown grouping function'):

cirq/work/observable_readout_calibration.py

@@ -1,6 +1,9 @@
 # pylint: disable=wrong-or-nonexistent-copyright-notice
+
+from __future__ import annotations
+
 import dataclasses
-from typing import Iterable, TYPE_CHECKING, Union
+from typing import Iterable, TYPE_CHECKING
 
 from cirq import circuits, ops, study
 from cirq.work.observable_measurement import measure_grouped_settings, StoppingCriteria
@@ -12,7 +15,7 @@ if TYPE_CHECKING:
 
 def calibrate_readout_error(
     qubits: Iterable[ops.Qid],
-    sampler: Union['cirq.Simulator', 'cirq.Sampler'],
+    sampler: cirq.Simulator | cirq.Sampler,
     stopping_criteria: StoppingCriteria,
 ):
     # We know there won't be any fancy sweeps or observables so we can