cirq-core 1.7.0.dev20250825174419__py3-none-any.whl → 1.7.0.dev20251203004401__py3-none-any.whl

cirq/sim/state_vector.py

@@ -17,7 +17,8 @@
 from __future__ import annotations
 
 import abc
-from typing import Mapping, Sequence, TYPE_CHECKING
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -105,7 +106,7 @@ class StateVectorMixin:
             and non-zero floats of the specified accuracy."""
         return qis.dirac_notation(self.state_vector(), decimals, qid_shape=self._qid_shape)
 
-    def density_matrix_of(self, qubits: list[cirq.Qid] | None = None) -> np.ndarray:
+    def density_matrix_of(self, qubits: Sequence[cirq.Qid] | None = None) -> np.ndarray:
         r"""Returns the density matrix of the state.
 
         Calculate the density matrix for the system on the qubits provided.
@@ -135,8 +136,7 @@ class StateVectorMixin:
 
         Raises:
             ValueError: if the size of the state represents more than 25 qubits.
-            IndexError: if the indices are out of range for the number of qubits
-                corresponding to the state.
+            KeyError: if some of the qubits provided are not in the quantum state.
         """
         return qis.density_matrix_from_state_vector(
             self.state_vector(),
@@ -160,8 +160,7 @@ class StateVectorMixin:
 
         Raises:
             ValueError: if the size of the state represents more than 25 qubits.
-            IndexError: if index is out of range for the number of qubits
-                corresponding to the state.
+            KeyError: if the specified qubit is not in the quantum state.
         """
         return qis.bloch_vector_from_state_vector(
             self.state_vector(), self.qubit_map[qubit], qid_shape=self._qid_shape

cirq/sim/state_vector_simulation_state.py

@@ -16,7 +16,8 @@
 
 from __future__ import annotations
 
-from typing import Any, Callable, Sequence, TYPE_CHECKING
+from collections.abc import Callable, Sequence
+from typing import Any, Self, TYPE_CHECKING
 
 import numpy as np
 
@@ -55,7 +56,7 @@ class _BufferedStateVector(qis.QuantumStateRepresentation):
         *,
         initial_state: np.ndarray | cirq.STATE_VECTOR_LIKE = 0,
         qid_shape: tuple[int, ...] | None = None,
-        dtype: type[np.complexfloating] | None = None,
+        dtype: type[np.complexfloating] | np.dtype[np.complexfloating] | None = None,
         buffer: np.ndarray | None = None,
     ):
         """Initializes the object with the inputs.
@@ -324,7 +325,7 @@ class StateVectorSimulationState(SimulationState[_BufferedStateVector]):
         prng: np.random.RandomState | None = None,
         qubits: Sequence[cirq.Qid] | None = None,
         initial_state: np.ndarray | cirq.STATE_VECTOR_LIKE = 0,
-        dtype: type[np.complexfloating] = np.complex64,
+        dtype: type[np.complexfloating] | np.dtype[np.complexfloating] = np.complex64,
         classical_data: cirq.ClassicalDataStore | None = None,
     ):
         """Inits StateVectorSimulationState.
@@ -356,7 +357,7 @@ class StateVectorSimulationState(SimulationState[_BufferedStateVector]):
         )
         super().__init__(state=state, prng=prng, qubits=qubits, classical_data=classical_data)
 
-    def add_qubits(self, qubits: Sequence[cirq.Qid]):
+    def add_qubits(self, qubits: Sequence[cirq.Qid]) -> Self:
         ret = super().add_qubits(qubits)
         return (
             self.kronecker_product(type(self)(qubits=qubits), inplace=True)
@@ -364,7 +365,7 @@ class StateVectorSimulationState(SimulationState[_BufferedStateVector]):
             else ret
         )
 
-    def remove_qubits(self, qubits: Sequence[cirq.Qid]):
+    def remove_qubits(self, qubits: Sequence[cirq.Qid]) -> Self:
         ret = super().remove_qubits(qubits)
         if ret is not NotImplemented:
             return ret
@@ -406,11 +407,11 @@ class StateVectorSimulationState(SimulationState[_BufferedStateVector]):
         )
 
     @property
-    def target_tensor(self):
+    def target_tensor(self) -> np.ndarray:
         return self._state._state_vector
 
     @property
-    def available_buffer(self):
+    def available_buffer(self) -> np.ndarray:
         return self._state._buffer
 
 
@@ -429,7 +430,7 @@ def _strat_act_on_state_vector_from_mixture(
     if index is None:
         return NotImplemented
     if protocols.is_measurement(action):
-        key = protocols.measurement_key_name(action)
+        key = protocols.measurement_key_obj(action)
         args._classical_data.record_channel_measurement(key, index)
     return True
 
@@ -441,6 +442,6 @@ def _strat_act_on_state_vector_from_channel(
     if index is None:
         return NotImplemented
     if protocols.is_measurement(action):
-        key = protocols.measurement_key_name(action)
+        key = protocols.measurement_key_obj(action)
         args._classical_data.record_channel_measurement(key, index)
     return True

cirq/sim/state_vector_simulator.py

@@ -18,8 +18,9 @@ from __future__ import annotations
 
 import abc
 import warnings
+from collections.abc import Iterator, Sequence
 from functools import cached_property
-from typing import Any, Generic, Iterator, Sequence, TYPE_CHECKING, TypeVar
+from typing import Any, Generic, TYPE_CHECKING, TypeVar
 
 import numpy as np
 

cirq/sim/state_vector_simulator_test.py

@@ -20,7 +20,7 @@ import cirq
 import cirq.testing
 
 
-def test_state_vector_trial_result_repr():
+def test_state_vector_trial_result_repr() -> None:
     q0 = cirq.NamedQubit('a')
     final_simulator_state = cirq.StateVectorSimulationState(
         available_buffer=np.array([0, 1], dtype=np.complex64),
@@ -47,7 +47,7 @@ def test_state_vector_trial_result_repr():
     assert eval(expected_repr) == trial_result
 
 
-def test_state_vector_trial_result_equality():
+def test_state_vector_trial_result_equality() -> None:
     eq = cirq.testing.EqualsTester()
     final_simulator_state = cirq.StateVectorSimulationState(initial_state=np.array([]))
     eq.add_equality_group(
@@ -86,7 +86,7 @@ def test_state_vector_trial_result_equality():
     )
 
 
-def test_state_vector_trial_result_state_mixin():
+def test_state_vector_trial_result_state_mixin() -> None:
     qubits = cirq.LineQubit.range(2)
     final_simulator_state = cirq.StateVectorSimulationState(
         qubits=qubits, initial_state=np.array([0, 1, 0, 0])
@@ -103,7 +103,7 @@ def test_state_vector_trial_result_state_mixin():
     assert result.dirac_notation() == '|01⟩'
 
 
-def test_state_vector_trial_result_qid_shape():
+def test_state_vector_trial_result_qid_shape() -> None:
     final_simulator_state = cirq.StateVectorSimulationState(
         qubits=[cirq.NamedQubit('a')], initial_state=np.array([0, 1])
     )
@@ -125,7 +125,7 @@ def test_state_vector_trial_result_qid_shape():
     assert cirq.qid_shape(trial_result) == (3, 2)
 
 
-def test_state_vector_trial_state_vector_is_copy():
+def test_state_vector_trial_state_vector_is_copy() -> None:
     final_state_vector = np.array([0, 1], dtype=np.complex64)
     qubit_map = {cirq.NamedQubit('a'): 0}
     final_simulator_state = cirq.StateVectorSimulationState(
@@ -137,7 +137,7 @@ def test_state_vector_trial_state_vector_is_copy():
     assert trial_result.state_vector(copy=True) is not final_simulator_state.target_tensor
 
 
-def test_state_vector_trial_result_no_qubits():
+def test_state_vector_trial_result_no_qubits() -> None:
     initial_state_vector = np.array([1], dtype=np.complex64)
     initial_state = initial_state_vector.reshape((2,) * 0)  # reshape as tensor for 0 qubits
     final_simulator_state = cirq.StateVectorSimulationState(qubits=[], initial_state=initial_state)
@@ -149,7 +149,7 @@ def test_state_vector_trial_result_no_qubits():
     assert np.array_equal(state_vector, initial_state_vector)
 
 
-def test_str_big():
+def test_str_big() -> None:
     qs = cirq.LineQubit.range(10)
     final_simulator_state = cirq.StateVectorSimulationState(
         prng=np.random.RandomState(0),
@@ -161,7 +161,7 @@ def test_str_big():
     assert 'output vector: [0.03125+0.j 0.03125+0.j 0.03125+0.j ..' in str(result)
 
 
-def test_str_qudit():
+def test_str_qudit() -> None:
     qutrit = cirq.LineQid(0, dimension=3)
     final_simulator_state = cirq.StateVectorSimulationState(
         prng=np.random.RandomState(0),
@@ -183,7 +183,7 @@ def test_str_qudit():
     assert "|1⟩" in str(result)
 
 
-def test_pretty_print():
+def test_pretty_print() -> None:
     final_simulator_state = cirq.StateVectorSimulationState(
         available_buffer=np.array([1]),
         prng=np.random.RandomState(0),

cirq/sim/state_vector_test.py

@@ -17,14 +17,13 @@
 from __future__ import annotations
 
 import itertools
-from typing import Iterator
+from collections.abc import Iterator
 from unittest import mock
 
 import numpy as np
 import pytest
 
 import cirq
-import cirq.testing
 from cirq import linalg
 
 
@@ -35,7 +34,7 @@ def use_np_transpose(request) -> Iterator[bool]:
         yield value
 
 
-def test_state_mixin():
+def test_state_mixin() -> None:
     class TestClass(cirq.StateVectorMixin):
         def state_vector(self, copy: bool | None = None) -> np.ndarray:
             return np.array([0, 0, 1, 0])
@@ -60,7 +59,7 @@ def test_state_mixin():
         _ = TestClass({qubits[0]: -1, qubits[1]: 1})
 
 
-def test_sample_state_big_endian():
+def test_sample_state_big_endian() -> None:
     results = []
     for x in range(8):
         state = cirq.to_valid_state_vector(x, 3)
@@ -71,7 +70,7 @@ def test_sample_state_big_endian():
         np.testing.assert_equal(result, expected)
 
 
-def test_sample_state_partial_indices():
+def test_sample_state_partial_indices() -> None:
     for index in range(3):
         for x in range(8):
             state = cirq.to_valid_state_vector(x, 3)
@@ -80,14 +79,14 @@ def test_sample_state_partial_indices():
             )
 
 
-def test_sample_state_partial_indices_oder():
+def test_sample_state_partial_indices_oder() -> None:
     for x in range(8):
         state = cirq.to_valid_state_vector(x, 3)
         expected = [[bool(1 & (x >> 0)), bool(1 & (x >> 1))]]
         np.testing.assert_equal(cirq.sample_state_vector(state, [2, 1]), expected)
 
 
-def test_sample_state_partial_indices_all_orders():
+def test_sample_state_partial_indices_all_orders() -> None:
     for perm in itertools.permutations([0, 1, 2]):
         for x in range(8):
             state = cirq.to_valid_state_vector(x, 3)
@@ -95,7 +94,7 @@ def test_sample_state_partial_indices_all_orders():
             np.testing.assert_equal(cirq.sample_state_vector(state, perm), expected)
 
 
-def test_sample_state():
+def test_sample_state() -> None:
     state = np.zeros(8, dtype=np.complex64)
     state[0] = 1 / np.sqrt(2)
     state[2] = 1 / np.sqrt(2)
@@ -110,12 +109,12 @@ def test_sample_state():
         np.testing.assert_equal(cirq.sample_state_vector(state, [0]), [[False]])
 
 
-def test_sample_empty_state():
+def test_sample_empty_state() -> None:
     state = np.array([1.0])
     np.testing.assert_almost_equal(cirq.sample_state_vector(state, []), np.zeros(shape=(1, 0)))
 
 
-def test_sample_no_repetitions():
+def test_sample_no_repetitions() -> None:
     state = cirq.to_valid_state_vector(0, 3)
     np.testing.assert_almost_equal(
         cirq.sample_state_vector(state, [1], repetitions=0), np.zeros(shape=(0, 1))
@@ -125,7 +124,7 @@ def test_sample_no_repetitions():
     )
 
 
-def test_sample_state_repetitions():
+def test_sample_state_repetitions() -> None:
     for perm in itertools.permutations([0, 1, 2]):
         for x in range(8):
             state = cirq.to_valid_state_vector(x, 3)
@@ -135,7 +134,7 @@ def test_sample_state_repetitions():
             np.testing.assert_equal(result, expected)
 
 
-def test_sample_state_seed():
+def test_sample_state_seed() -> None:
     state = np.ones(2) / np.sqrt(2)
 
     samples = cirq.sample_state_vector(state, [0], repetitions=10, seed=1234)
@@ -151,20 +150,20 @@ def test_sample_state_seed():
     )
 
 
-def test_sample_state_negative_repetitions():
+def test_sample_state_negative_repetitions() -> None:
     state = cirq.to_valid_state_vector(0, 3)
     with pytest.raises(ValueError, match='-1'):
         cirq.sample_state_vector(state, [1], repetitions=-1)
 
 
-def test_sample_state_not_power_of_two():
+def test_sample_state_not_power_of_two() -> None:
     with pytest.raises(ValueError, match='3'):
         cirq.sample_state_vector(np.array([1, 0, 0]), [1])
     with pytest.raises(ValueError, match='5'):
         cirq.sample_state_vector(np.array([0, 1, 0, 0, 0]), [1])
 
 
-def test_sample_state_index_out_of_range():
+def test_sample_state_index_out_of_range() -> None:
     state = cirq.to_valid_state_vector(0, 3)
     with pytest.raises(IndexError, match='-2'):
         cirq.sample_state_vector(state, [-2])
@@ -172,12 +171,12 @@ def test_sample_state_index_out_of_range():
         cirq.sample_state_vector(state, [3])
 
 
-def test_sample_no_indices():
+def test_sample_no_indices() -> None:
     state = cirq.to_valid_state_vector(0, 3)
     np.testing.assert_almost_equal(cirq.sample_state_vector(state, []), np.zeros(shape=(1, 0)))
 
 
-def test_sample_no_indices_repetitions():
+def test_sample_no_indices_repetitions() -> None:
     state = cirq.to_valid_state_vector(0, 3)
     np.testing.assert_almost_equal(
         cirq.sample_state_vector(state, [], repetitions=2), np.zeros(shape=(2, 0))
@@ -185,7 +184,7 @@ def test_sample_no_indices_repetitions():
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_computational_basis(use_np_transpose: bool):
+def test_measure_state_computational_basis(use_np_transpose: bool) -> None:
     # verify patching of can_numpy_support_shape in the use_np_transpose fixture
     assert linalg.can_numpy_support_shape([1]) is use_np_transpose
     results = []
@@ -199,7 +198,7 @@ def test_measure_state_computational_basis(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_reshape(use_np_transpose: bool):
+def test_measure_state_reshape(use_np_transpose: bool) -> None:
     results = []
     for x in range(8):
         initial_state = np.reshape(cirq.to_valid_state_vector(x, 3), [2] * 3)
@@ -211,7 +210,7 @@ def test_measure_state_reshape(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_partial_indices(use_np_transpose: bool):
+def test_measure_state_partial_indices(use_np_transpose: bool) -> None:
     for index in range(3):
         for x in range(8):
             initial_state = cirq.to_valid_state_vector(x, 3)
@@ -221,7 +220,7 @@ def test_measure_state_partial_indices(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_partial_indices_order(use_np_transpose: bool):
+def test_measure_state_partial_indices_order(use_np_transpose: bool) -> None:
     for x in range(8):
         initial_state = cirq.to_valid_state_vector(x, 3)
         bits, state = cirq.measure_state_vector(initial_state, [2, 1])
@@ -230,7 +229,7 @@ def test_measure_state_partial_indices_order(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_partial_indices_all_orders(use_np_transpose: bool):
+def test_measure_state_partial_indices_all_orders(use_np_transpose: bool) -> None:
     for perm in itertools.permutations([0, 1, 2]):
         for x in range(8):
             initial_state = cirq.to_valid_state_vector(x, 3)
@@ -240,7 +239,7 @@ def test_measure_state_partial_indices_all_orders(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_collapse(use_np_transpose: bool):
+def test_measure_state_collapse(use_np_transpose: bool) -> None:
     initial_state = np.zeros(8, dtype=np.complex64)
     initial_state[0] = 1 / np.sqrt(2)
     initial_state[2] = 1 / np.sqrt(2)
@@ -264,7 +263,7 @@ def test_measure_state_collapse(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_seed(use_np_transpose: bool):
+def test_measure_state_seed(use_np_transpose: bool) -> None:
     n = 10
     initial_state = np.ones(2**n) / 2 ** (n / 2)
 
@@ -284,7 +283,7 @@ def test_measure_state_seed(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_out_is_state(use_np_transpose: bool):
+def test_measure_state_out_is_state(use_np_transpose: bool) -> None:
     initial_state = np.zeros(8, dtype=np.complex64)
     initial_state[0] = 1 / np.sqrt(2)
     initial_state[2] = 1 / np.sqrt(2)
@@ -296,7 +295,7 @@ def test_measure_state_out_is_state(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_out_is_not_state(use_np_transpose: bool):
+def test_measure_state_out_is_not_state(use_np_transpose: bool) -> None:
     initial_state = np.zeros(8, dtype=np.complex64)
     initial_state[0] = 1 / np.sqrt(2)
     initial_state[2] = 1 / np.sqrt(2)
@@ -307,7 +306,7 @@ def test_measure_state_out_is_not_state(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_not_power_of_two(use_np_transpose: bool):
+def test_measure_state_not_power_of_two(use_np_transpose: bool) -> None:
     with pytest.raises(ValueError, match='3'):
         _, _ = cirq.measure_state_vector(np.array([1, 0, 0]), [1])
     with pytest.raises(ValueError, match='5'):
@@ -315,7 +314,7 @@ def test_measure_state_not_power_of_two(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_index_out_of_range(use_np_transpose: bool):
+def test_measure_state_index_out_of_range(use_np_transpose: bool) -> None:
     state = cirq.to_valid_state_vector(0, 3)
     with pytest.raises(IndexError, match='-2'):
         cirq.measure_state_vector(state, [-2])
@@ -324,7 +323,7 @@ def test_measure_state_index_out_of_range(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_no_indices(use_np_transpose: bool):
+def test_measure_state_no_indices(use_np_transpose: bool) -> None:
     initial_state = cirq.to_valid_state_vector(0, 3)
     bits, state = cirq.measure_state_vector(initial_state, [])
     assert [] == bits
@@ -332,7 +331,7 @@ def test_measure_state_no_indices(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_no_indices_out_is_state(use_np_transpose: bool):
+def test_measure_state_no_indices_out_is_state(use_np_transpose: bool) -> None:
     initial_state = cirq.to_valid_state_vector(0, 3)
     bits, state = cirq.measure_state_vector(initial_state, [], out=initial_state)
     assert [] == bits
@@ -341,7 +340,7 @@ def test_measure_state_no_indices_out_is_state(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_no_indices_out_is_not_state(use_np_transpose: bool):
+def test_measure_state_no_indices_out_is_not_state(use_np_transpose: bool) -> None:
     initial_state = cirq.to_valid_state_vector(0, 3)
     out = np.zeros_like(initial_state)
     bits, state = cirq.measure_state_vector(initial_state, [], out=out)
@@ -352,7 +351,7 @@ def test_measure_state_no_indices_out_is_not_state(use_np_transpose: bool):
 
 
 @pytest.mark.parametrize('use_np_transpose', [False, True], indirect=True)
-def test_measure_state_empty_state(use_np_transpose: bool):
+def test_measure_state_empty_state(use_np_transpose: bool) -> None:
     initial_state = np.array([1.0])
     bits, state = cirq.measure_state_vector(initial_state, [])
     assert [] == bits
@@ -364,12 +363,12 @@ class BasicStateVector(cirq.StateVectorMixin):
         return np.array([0, 1, 0, 0])
 
 
-def test_step_result_pretty_state():
+def test_step_result_pretty_state() -> None:
     step_result = BasicStateVector()
     assert step_result.dirac_notation() == '|01⟩'
 
 
-def test_step_result_density_matrix():
+def test_step_result_density_matrix() -> None:
     q0, q1 = cirq.LineQubit.range(2)
 
     step_result = BasicStateVector({q0: 0, q1: 1})
@@ -385,7 +384,7 @@ def test_step_result_density_matrix():
     np.testing.assert_array_almost_equal(single_rho, step_result.density_matrix_of([q1]))
 
 
-def test_step_result_density_matrix_invalid():
+def test_step_result_density_matrix_invalid() -> None:
     q0, q1 = cirq.LineQubit.range(2)
 
     step_result = BasicStateVector({q0: 0})
@@ -393,21 +392,23 @@ def test_step_result_density_matrix_invalid():
     with pytest.raises(KeyError):
         step_result.density_matrix_of([q1])
     with pytest.raises(KeyError):
-        step_result.density_matrix_of('junk')
+        step_result.density_matrix_of('junk')  # type: ignore[arg-type]
     with pytest.raises(TypeError):
-        step_result.density_matrix_of(0)
+        step_result.density_matrix_of(0)  # type: ignore[arg-type]
 
 
-def test_step_result_bloch_vector():
+def test_step_result_bloch_vector() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     step_result = BasicStateVector({q0: 0, q1: 1})
     bloch1 = np.array([0, 0, -1])
     bloch0 = np.array([0, 0, 1])
     np.testing.assert_array_almost_equal(bloch1, step_result.bloch_vector_of(q1))
     np.testing.assert_array_almost_equal(bloch0, step_result.bloch_vector_of(q0))
+    with pytest.raises(KeyError):
+        step_result.bloch_vector_of(cirq.LineQubit(2))
 
 
-def test_factor_validation():
+def test_factor_validation() -> None:
     args = cirq.Simulator()._create_simulation_state(0, qubits=cirq.LineQubit.range(2))
     args.apply_operation(cirq.H(cirq.LineQubit(0)) ** 0.7)
     t = args.create_merged_state().target_tensor

cirq/study/__init__.py

@@ -38,6 +38,7 @@ from cirq.study.sweepable import (
 from cirq.study.sweeps import (
     Concat as Concat,
     Linspace as Linspace,
+    list_of_dicts_to_zip as list_of_dicts_to_zip,
     ListSweep as ListSweep,
     Points as Points,
     Product as Product,

cirq/study/flatten_expressions.py

@@ -17,7 +17,8 @@
 from __future__ import annotations
 
 import numbers
-from typing import Any, Callable, TYPE_CHECKING
+from collections.abc import Callable
+from typing import Any, TYPE_CHECKING
 
 import sympy
 

cirq/study/resolver.py

@@ -17,7 +17,8 @@
 from __future__ import annotations
 
 import numbers
-from typing import Any, cast, Iterator, Mapping, TYPE_CHECKING, Union
+from collections.abc import Iterator, Mapping
+from typing import Any, cast, TYPE_CHECKING, Union
 
 import numpy as np
 import sympy
@@ -75,9 +76,20 @@ class ParamResolver:
 
         self._param_hash: int | None = None
         self._param_dict = cast(ParamDictType, {} if param_dict is None else param_dict)
+        self._param_dict_with_str_keys = self._param_dict
+        generate_str_keys = False
         for key in self._param_dict:
-            if isinstance(key, sympy.Expr) and not isinstance(key, sympy.Symbol):
-                raise TypeError(f'ParamResolver keys cannot be (non-symbol) formulas ({key})')
+            if isinstance(key, sympy.Expr):
+                if isinstance(key, sympy.Symbol):
+                    generate_str_keys = True
+                else:
+                    raise TypeError(f'ParamResolver keys cannot be (non-symbol) formulas ({key})')
+        if generate_str_keys:
+            # Remake dictionary with string keys for faster access
+            self._param_dict_with_str_keys = {
+                (key.name if isinstance(key, sympy.Symbol) else key): value
+                for key, value in self._param_dict.items()
+            }
         self._deep_eval_map: ParamDictType = {}
 
     @property
@@ -118,32 +130,33 @@ class ParamResolver:
             sympy.SympifyError: If the resulting value cannot be interpreted.
         """
 
-        # Input is a pass through type, no resolution needed: return early
-        v = _resolve_value(value)
-        if v is not NotImplemented:
-            return v
-
         # Handle string or symbol
-        if isinstance(value, (str, sympy.Symbol)):
-            string = value if isinstance(value, str) else value.name
-            symbol = value if isinstance(value, sympy.Symbol) else sympy.Symbol(value)
-            param_value = self._param_dict.get(string, _NOT_FOUND)
-            if param_value is _NOT_FOUND:
-                param_value = self._param_dict.get(symbol, _NOT_FOUND)
+        original_value = value
+        if isinstance(value, sympy.Symbol):
+            value = value.name
+        if isinstance(value, str):
+            param_value = self._param_dict_with_str_keys.get(value, _NOT_FOUND)
+            if isinstance(param_value, float):
+                return param_value
             if param_value is _NOT_FOUND:
                 # Symbol or string cannot be resolved if not in param dict; return as symbol.
-                return symbol
+                return sympy.Symbol(value)
             v = _resolve_value(param_value)
             if v is not NotImplemented:
                 return v
             if isinstance(param_value, str):
                 param_value = sympy.Symbol(param_value)
             elif not isinstance(param_value, sympy.Basic):
-                return value
+                return original_value
             if recursive:
                 param_value = self._value_of_recursive(value)
             return param_value
 
+        # Input is a pass through type, no resolution needed: return early
+        v = _resolve_value(value)
+        if v is not NotImplemented:
+            return v
+
         if not isinstance(value, sympy.Basic):
             # No known way to resolve this variable, return unchanged.
             return value
@@ -210,7 +223,7 @@ class ParamResolver:
         self._deep_eval_map[value] = _RECURSION_FLAG
 
         v = self.value_of(value, recursive=False)
-        if v == value:
+        if v == value or (isinstance(v, sympy.Symbol) and v.name == value):
             self._deep_eval_map[value] = v
         else:
             self._deep_eval_map[value] = self.value_of(v, recursive=True)
@@ -278,7 +291,7 @@ class ParamResolver:
 
 
 def _resolve_value(val: Any) -> Any:
-    if val is None:
+    if isinstance(val, float) or val is None:
         return val
     if isinstance(val, numbers.Number) and not isinstance(val, sympy.Basic):
         return val

cirq/study/resolver_test.py

@@ -31,7 +31,7 @@ import cirq
         None,
         3.2,
         np.float32(3.2),
-        int(1),
+        1,
         np.int32(45),
         np.float64(6.3),
         np.int32(2),

cirq/study/result.py

@@ -19,7 +19,8 @@ from __future__ import annotations
 import abc
 import collections
 import io
-from typing import Any, Callable, cast, Iterable, Mapping, Sequence, TYPE_CHECKING, TypeVar, Union
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from typing import Any, cast, TYPE_CHECKING, TypeVar, Union
 
 import numpy as np
 import pandas as pd