More numpy types

cirq-core/cirq/circuits/circuit.py

@@ -1592,7 +1592,7 @@ def _control_keys_(self) -> FrozenSet['cirq.MeasurementKey']:
 
 
 def _overlap_collision_time(
-    c1: Sequence['cirq.Moment'], c2: Sequence['cirq.Moment'], align: 'cirq.Alignment'
+    c1: np.ndarray, c2: Sequence['cirq.Moment'], align: 'cirq.Alignment'
 ) -> int:
     # Tracks the first used moment index for each qubit in c2.
     # Tracks the complementary last used moment index for each qubit in c1.

cirq-core/cirq/circuits/circuit_operation.py

@@ -589,7 +589,7 @@ def repeat(
                 # As CircuitOperation is immutable, this can safely return the original.
                 return self
 
-            expected_repetition_id_length = abs(repetitions)
+            expected_repetition_id_length: int = np.abs(repetitions)
 
             if repetition_ids is None:
                 if self.use_repetition_ids:

cirq-core/cirq/experiments/cross_entropy_benchmarking.py

@@ -12,7 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Mapping, NamedTuple, Optional, Sequence, TYPE_CHECKING, Tuple
+from typing import (
+    Any,
+    Dict,
+    List,
+    Mapping,
+    NamedTuple,
+    Optional,
+    Sequence,
+    TYPE_CHECKING,
+    Tuple,
+    Union,
+)
 import dataclasses
 import numpy as np
 from matplotlib import pyplot as plt
@@ -208,7 +219,9 @@ def __repr__(self) -> str:
         return f'cirq.experiments.CrossEntropyResult({args})'
 
 
-def _fit_exponential_decay(x: Sequence[int], y: Sequence[float]) -> Tuple[np.ndarray, np.ndarray]:
+def _fit_exponential_decay(
+    x: Union[Sequence[int], np.ndarray], y: Union[Sequence[float], np.ndarray]
+) -> Tuple[np.ndarray, np.ndarray]:
     """Fit an exponential model y = S * p**x using nonlinear least squares.
 
     Args:

cirq-core/cirq/ops/eigen_gate.py

@@ -331,6 +331,7 @@ def _trace_distance_bound_(self) -> Optional[float]:
         if protocols.is_parameterized(self._exponent):
             return None
         angles = np.pi * (np.array(self._eigen_shifts()) * self._exponent % 2)
+        print(type(angles))
         return protocols.trace_distance_from_angle_list(angles)
 
     def _has_unitary_(self) -> bool:

cirq-core/cirq/ops/pauli_string.py

@@ -644,7 +644,7 @@ def expectation_from_density_matrix(
         *,
         atol: float = 1e-7,
         check_preconditions: bool = True,
-    ) -> float:
+    ) -> np.ndarray:
         r"""Evaluate the expectation of this PauliString given a density matrix.
 
         Compute the expectation value of this PauliString with respect to an
@@ -716,7 +716,7 @@ def expectation_from_density_matrix(
 
     def _expectation_from_density_matrix_no_validation(
         self, state: np.ndarray, qubit_map: Mapping[TKey, int]
-    ) -> float:
+    ) -> np.ndarray:
         """Evaluate the expectation of this PauliString given a density matrix.
 
         This method does not provide input validation. See

cirq-core/cirq/protocols/circuit_diagram_info_protocol.py

@@ -274,7 +274,7 @@ def format_radians(self, radians: Union[sympy.Basic, int, float]) -> str:
             return '0'
         if radians == -np.pi:
             return '-' + unit
-        if self.precision is not None:
+        if self.precision is not None and not isinstance(radians, sympy.Basic):
             quantity = self.format_real(radians / np.pi)
             return quantity + unit
         return repr(radians)

cirq-core/cirq/protocols/trace_distance_bound.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, TypeVar, Optional, Sequence
+from typing import Any, TypeVar, Optional, Sequence, Union
 
 import numpy as np
 from typing_extensions import Protocol
@@ -109,7 +109,7 @@ def _strat_distance_from_unitary(val: Any) -> Optional[float]:
     return trace_distance_from_angle_list(np.angle(np.linalg.eigvals(u)))
 
 
-def trace_distance_from_angle_list(angle_list: Sequence[float]) -> float:
+def trace_distance_from_angle_list(angle_list: Union[Sequence[float], np.ndarray]) -> float:
     """Given a list of arguments of the eigenvalues of a unitary matrix,
     calculates the trace distance bound of the unitary effect.
 

cirq-core/cirq/qis/clifford_tableau.py

@@ -19,7 +19,7 @@
 from cirq import protocols
 from cirq._compat import proper_repr
 from cirq.qis import quantum_state_representation
-from cirq.value import big_endian_int_to_digits, linear_dict
+from cirq.value import big_endian_int_to_digits, linear_dict, random_state
 
 if TYPE_CHECKING:
     import cirq
@@ -509,11 +509,14 @@ def destabilizers(self) -> List['cirq.DensePauliString']:
         generators above generate the full Pauli group on n qubits."""
         return [self._row_to_dense_pauli(i) for i in range(self.n)]
 
-    def _measure(self, q, prng: np.random.RandomState = np.random) -> int:
+    def _measure(self, q, prng: Optional[np.random.RandomState] = None) -> int:
         """Performs a projective measurement on the q'th qubit.
 
         Returns: the result (0 or 1) of the measurement.
         """
+        real_prng: np.random.RandomState = (
+            random_state.parse_random_state(np.random) if prng is None else prng
+        )
         is_commuting = True
         for i in range(self.n, 2 * self.n):
             if self.xs[i, q]:
@@ -544,7 +547,7 @@ def _measure(self, q, prng: np.random.RandomState = np.random) -> int:
 
         self.zs[p, q] = True
 
-        self.rs[p] = bool(prng.randint(2))
+        self.rs[p] = bool(real_prng.randint(2))
 
         return int(self.rs[p])
 

cirq-core/cirq/sim/clifford/stabilizer_state_ch_form.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Sequence, Union
+from typing import Any, Dict, List, Sequence, Optional, Union
 import numpy as np
 
 import cirq
 from cirq import protocols, qis, value
-from cirq.value import big_endian_int_to_digits
+from cirq.value import big_endian_int_to_digits, random_state
 
 
 @value.value_equality
@@ -236,17 +236,20 @@ def to_state_vector(self) -> np.ndarray:
 
         return arr
 
-    def _measure(self, q, prng: np.random.RandomState) -> int:
+    def _measure(self, q, prng: Optional[np.random.RandomState] = None) -> int:
         """Measures the q'th qubit.
 
         Reference: Section 4.1 "Simulating measurements"
 
         Returns: Computational basis measurement as 0 or 1.
         """
+        real_prng: np.random.RandomState = (
+            random_state.parse_random_state(np.random) if prng is None else prng
+        )
         w = self.s.copy()
         for i, v_i in enumerate(self.v):
             if v_i == 1:
-                w[i] = bool(prng.randint(2))
+                w[i] = bool(real_prng.randint(2))
         x_i = sum(w & self.G[q, :]) % 2
         # Project the state to the above measurement outcome.
         self.project_Z(q, x_i)
@@ -386,7 +389,7 @@ def apply_global_phase(self, coefficient: value.Scalar):
         self.omega *= coefficient
 
     def measure(
-        self, axes: Sequence[int], seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None
+        self, axes: Sequence[int], seed: Optional['cirq.RANDOM_STATE_OR_SEED_LIKE'] = None
     ) -> List[int]:
         return [self._measure(axis, seed) for axis in axes]
 

cirq-core/cirq/sim/density_matrix_utils.py

@@ -68,7 +68,6 @@ def sample_density_matrix(
         qid_shape = (2,) * num_qubits
     else:
         _validate_density_matrix_qid_shape(density_matrix, qid_shape)
-        num_qubits = len(qid_shape)
     meas_shape = _indices_shape(qid_shape, indices)
 
     if repetitions == 0 or len(indices) == 0:
@@ -139,15 +138,18 @@ def measure_density_matrix(
         qid_shape = (2,) * num_qubits
     else:
         _validate_density_matrix_qid_shape(density_matrix, qid_shape)
-        num_qubits = len(qid_shape)
     meas_shape = _indices_shape(qid_shape, indices)
 
+    def _copy_density_matrix_to_out(density_matrix: np.ndarray, out: np.ndarray) -> np.ndarray:
+        np.copyto(dst=out, src=density_matrix)
+        return out
+
     arrout: np.ndarray = (
         np.copy(density_matrix)
         if out is None
         else density_matrix
         if out is density_matrix
-        else (np.copyto(dst=out, src=density_matrix), out)[-1]
+        else _copy_density_matrix_to_out(density_matrix, out)
     )
 
     if len(indices) == 0:

cirq-core/cirq/sim/state_vector_simulation_state.py

@@ -226,7 +226,7 @@ def prepare_into_buffer(k: int):
         ]
         p = prng.random()
         weight = None
-        fallback_weight = 0
+        fallback_weight = np.float64(0)
         fallback_weight_index = 0
         index = None
         for index in range(len(kraus_tensors)):

cirq-core/cirq/transformers/heuristic_decompositions/gate_tabulation_math_utils.py

@@ -3,6 +3,7 @@
 from typing import Union, Sequence, Optional
 
 import numpy as np
+from cirq.value import random_state
 
 _RealArraylike = Union[np.ndarray, float]
 
@@ -58,7 +59,9 @@ def random_qubit_unitary(
         rng: Random number generator to be used in sampling. Default is
             numpy.random.
     """
-    real_rng: np.random.RandomState = np.random if rng is None else rng
+    real_rng: np.random.RandomState = (
+        random_state.parse_random_state(np.random) if rng is None else rng
+    )
 
     theta = np.arcsin(np.sqrt(real_rng.rand(*shape)))
     phi_d = real_rng.rand(*shape) * np.pi * 2

cirq-core/cirq/value/type_alias.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from typing import Union
+
 import sympy
 
 from cirq._doc import document

cirq-core/cirq/vis/state_histogram.py

@@ -93,9 +93,12 @@ def plot_state_histogram(
         tick_label, values = zip(*sorted(data.items()))
     else:
         values = np.array(data)
-    if not tick_label:
-        tick_label = np.arange(len(values))
-    ax.bar(np.arange(len(values)), values, tick_label=tick_label)
+    values_range = np.arange(len(values))
+    _ = (
+        ax.bar(values_range, values, tick_label=values_range)
+        if tick_label is None
+        else ax.bar(values_range, values, tick_label=tick_label)
+    )
     ax.set_xlabel(xlabel)
     ax.set_ylabel(ylabel)
     ax.set_title(title)