Fix more mypy --next type errors (quantumlib#5392)

cirq/study/resolver.py

@@ -157,9 +157,11 @@ def value_of(
         if isinstance(value, sympy.Pow) and len(value.args) == 2:
             base = self.value_of(value.args[0], recursive)
             exponent = self.value_of(value.args[1], recursive)
+            # Casts because numpy can handle expressions (by delegating to __pow__), but does
+            # not have signature that will support this.
             if isinstance(base, numbers.Number):
-                return np.float_power(base, exponent)
-            return np.power(base, exponent)
+                return np.float_power(cast(complex, base), cast(complex, exponent))
+            return np.power(cast(complex, base), cast(complex, exponent))
 
         if not isinstance(value, sympy.Basic):
             # No known way to resolve this variable, return unchanged.

cirq/transformers/eject_phased_paulis.py

@@ -321,7 +321,7 @@ def _try_get_known_phased_pauli(
     elif (
         isinstance(gate, ops.PhasedXZGate)
         and not protocols.is_parameterized(gate.z_exponent)
-        and np.isclose(gate.z_exponent, 0)
+        and np.isclose(float(gate.z_exponent), 0)
     ):
         e = gate.x_exponent
         p = gate.axis_phase_exponent
@@ -336,9 +336,12 @@ def _try_get_known_z_half_turns(
     g = op.gate
     if (
         isinstance(g, ops.PhasedXZGate)
-        and np.isclose(g.x_exponent, 0)
-        and np.isclose(g.axis_phase_exponent, 0)
+        and not protocols.is_parameterized(g.x_exponent)
+        and not protocols.is_parameterized(g.axis_phase_exponent)
+        and np.isclose(float(g.x_exponent), 0)
+        and np.isclose(float(g.axis_phase_exponent), 0)
     ):
+
         h = g.z_exponent
     elif isinstance(g, ops.ZPowGate):
         h = g.exponent

cirq/transformers/heuristic_decompositions/gate_tabulation_math_utils.py

@@ -58,17 +58,17 @@ def random_qubit_unitary(
         rng: Random number generator to be used in sampling. Default is
             numpy.random.
     """
-    rng = np.random if rng is None else rng
+    real_rng: np.random.RandomState = np.random if rng is None else rng
 
-    theta = np.arcsin(np.sqrt(rng.rand(*shape)))
-    phi_d = rng.rand(*shape) * np.pi * 2
-    phi_o = rng.rand(*shape) * np.pi * 2
+    theta = np.arcsin(np.sqrt(real_rng.rand(*shape)))
+    phi_d = real_rng.rand(*shape) * np.pi * 2
+    phi_o = real_rng.rand(*shape) * np.pi * 2
 
     out = _single_qubit_unitary(theta, phi_d, phi_o)
 
     if randomize_global_phase:
         out = np.moveaxis(out, (-2, -1), (0, 1))
-        out *= np.exp(1j * np.pi * 2 * rng.rand(*shape))
+        out *= np.exp(1j * np.pi * 2 * real_rng.rand(*shape))
         out = np.moveaxis(out, (0, 1), (-2, -1))
     return out
 

cirq/transformers/heuristic_decompositions/two_qubit_gate_tabulation.py

@@ -15,7 +15,7 @@
 """Attempt to tabulate single qubit gates required to generate a target 2Q gate
 with a product A k A."""
 from functools import reduce
-from typing import Tuple, Sequence, List, NamedTuple
+from typing import List, NamedTuple, Sequence, Tuple
 
 from dataclasses import dataclass
 import numpy as np
@@ -100,7 +100,7 @@ def compile_two_qubit_gate(self, unitary: np.ndarray) -> TwoQubitGateTabulationR
         unitary = np.asarray(unitary)
         kak_vec = cirq.kak_vector(unitary, check_preconditions=False)
         infidelities = kak_vector_infidelity(kak_vec, self.kak_vecs, ignore_equivalent_vectors=True)
-        nearest_ind = infidelities.argmin()
+        nearest_ind = int(infidelities.argmin())
 
         success = infidelities[nearest_ind] < self.max_expected_infidelity
 
@@ -483,13 +483,13 @@ def two_qubit_gate_product_tabulation(
         else:
             missed_points.append(missing_vec)
 
-    kak_vecs = np.array(kak_vecs)
+    kak_vecs_arr = np.array(kak_vecs)
     summary += (
         f'\nFraction of Weyl chamber reached with 2 gates and 3 gates '
         f'(after patchup)'
-        f': {(len(kak_vecs) - 1) / num_mesh_points :.3f}'
+        f': {(len(kak_vecs_arr) - 1) / num_mesh_points :.3f}'
     )
 
     return TwoQubitGateTabulation(
-        base_gate, kak_vecs, sq_cycles, max_infidelity, summary, tuple(missed_points)
+        base_gate, kak_vecs_arr, sq_cycles, max_infidelity, summary, tuple(missed_points)
     )

cirq/value/duration.py

@@ -88,11 +88,8 @@ def __init__(
             else:
                 raise TypeError(f'Not a `cirq.DURATION_LIKE`: {repr(value)}.')
 
-        self._picos: Union[float, int, sympy.Expr] = (
-            picos + nanos * 1000 + micros * 1000_000 + millis * 1000_000_000
-        )
-        if isinstance(self._picos, np.number):
-            self._picos = float(self._picos)
+        val = picos + nanos * 1000 + micros * 1000_000 + millis * 1000_000_000
+        self._picos: _NUMERIC_OUTPUT_TYPE = float(val) if isinstance(val, np.number) else val
 
     def _is_parameterized_(self) -> bool:
         return protocols.is_parameterized(self._picos)

cirq/vis/histogram.py

@@ -86,15 +86,16 @@ def integrated_histogram(
     if isinstance(data, Mapping):
         data = list(data.values())
 
-    data = [d for d in data if not np.isnan(d)]
-    n = len(data)
+    float_data = [float(d) for d in data if not np.isnan(float(d))]
+
+    n = len(float_data)
 
     if not show_zero:
         bin_values = np.linspace(0, 1, n + 1)
-        parameter_values = sorted(np.concatenate(([0], data)))
+        parameter_values = sorted(np.concatenate(([0], float_data)))
     else:
         bin_values = np.linspace(0, 1, n)
-        parameter_values = sorted(data)
+        parameter_values = sorted(float_data)
     plot_options = {"where": 'post', "color": 'b', "linestyle": '-', "lw": 1.0, "ms": 0.0}
     plot_options.update(kwargs)
 
@@ -127,15 +128,19 @@ def integrated_histogram(
 
     if median_line:
         set_line(
-            np.median(data),
+            np.median(float_data),
             linestyle='--',
             color=plot_options['color'],
             alpha=0.5,
             label=median_label,
         )
     if mean_line:
         set_line(
-            np.mean(data), linestyle='-.', color=plot_options['color'], alpha=0.5, label=mean_label
+            np.mean(float_data),
+            linestyle='-.',
+            color=plot_options['color'],
+            alpha=0.5,
+            label=mean_label,
         )
     if show_plot:
         fig.show()

cirq/vis/state_histogram.py

@@ -92,7 +92,7 @@ def plot_state_histogram(
     elif isinstance(data, collections.Counter):
         tick_label, values = zip(*sorted(data.items()))
     else:
-        values = data
+        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)