Fix most numpy type errors in cirq/ops (#3997)

Using `check/mypy --next | grep cirq/ops` this fixes almost all of the issues raised with simple modifications.

There is a looming problem with places where we use `numbers.Complex`, which is a nice generalization for Union[int, float, complex] but does not play nice with numpy type information.

Towards #3767

Author: mpharrigan

cirq/ops/common_channels.py

@@ -194,10 +194,10 @@ def _json_dict_(self) -> Dict[str, Any]:
     def _approx_eq_(self, other: Any, atol: float) -> bool:
         return (
             self.num_qubits == other.num_qubits
-            and np.isclose(self.p_i, other.p_i, atol=atol)
-            and np.isclose(self.p_x, other.p_x, atol=atol)
-            and np.isclose(self.p_y, other.p_y, atol=atol)
-            and np.isclose(self.p_z, other.p_z, atol=atol)
+            and np.isclose(self.p_i, other.p_i, atol=atol).item()
+            and np.isclose(self.p_x, other.p_x, atol=atol).item()
+            and np.isclose(self.p_y, other.p_y, atol=atol).item()
+            and np.isclose(self.p_z, other.p_z, atol=atol).item()
         )
 
 
@@ -356,7 +356,7 @@ def _json_dict_(self) -> Dict[str, Any]:
         return protocols.obj_to_dict_helper(self, ['p', 'n_qubits'])
 
     def _approx_eq_(self, other: Any, atol: float) -> bool:
-        return np.isclose(self.p, other.p, atol=atol) and self.n_qubits == other.n_qubits
+        return np.isclose(self.p, other.p, atol=atol).item() and self.n_qubits == other.n_qubits
 
 
 def depolarize(p: float, n_qubits: int = 1) -> DepolarizingChannel:
@@ -498,8 +498,9 @@ def _json_dict_(self) -> Dict[str, Any]:
         return protocols.obj_to_dict_helper(self, ['p', 'gamma'])
 
     def _approx_eq_(self, other: Any, atol: float) -> bool:
-        return np.isclose(self.gamma, other.gamma, atol=atol) and np.isclose(
-            self.p, other.p, atol=atol
+        return (
+            np.isclose(self.gamma, other.gamma, atol=atol).item()
+            and np.isclose(self.p, other.p, atol=atol).item()
         )
 
 
@@ -628,7 +629,7 @@ def _json_dict_(self) -> Dict[str, Any]:
         return protocols.obj_to_dict_helper(self, ['gamma'])
 
     def _approx_eq_(self, other: Any, atol: float) -> bool:
-        return np.isclose(self.gamma, other.gamma, atol=atol)
+        return np.isclose(self.gamma, other.gamma, atol=atol).item()
 
 
 def amplitude_damp(gamma: float) -> AmplitudeDampingChannel:
@@ -863,7 +864,7 @@ def _json_dict_(self) -> Dict[str, Any]:
         return protocols.obj_to_dict_helper(self, ['gamma'])
 
     def _approx_eq_(self, other: Any, atol: float) -> bool:
-        return np.isclose(self._gamma, other._gamma, atol=atol)
+        return np.isclose(self._gamma, other._gamma, atol=atol).item()
 
 
 def phase_damp(gamma: float) -> PhaseDampingChannel:
@@ -973,7 +974,7 @@ def _json_dict_(self) -> Dict[str, Any]:
         return protocols.obj_to_dict_helper(self, ['p'])
 
     def _approx_eq_(self, other: Any, atol: float) -> bool:
-        return np.isclose(self.p, other.p, atol=atol)
+        return np.isclose(self.p, other.p, atol=atol).item()
 
 
 def _phase_flip_Z() -> common_gates.ZPowGate:
@@ -1129,7 +1130,7 @@ def _json_dict_(self) -> Dict[str, Any]:
         return protocols.obj_to_dict_helper(self, ['p'])
 
     def _approx_eq_(self, other: Any, atol: float) -> bool:
-        return np.isclose(self._p, other._p, atol=atol)
+        return np.isclose(self._p, other._p, atol=atol).item()
 
 
 def _bit_flip(p: float) -> BitFlipChannel:

cirq/ops/controlled_operation.py

@@ -137,7 +137,7 @@ def _extend_matrix(self, sub_matrix: np.ndarray) -> np.ndarray:
         for control_vals in itertools.product(*self.control_values):
             active = (*(v for v in control_vals), *(slice(None),) * sub_n) * 2
             tensor[active] = sub_tensor
-        return tensor.reshape((np.prod(qid_shape, dtype=int),) * 2)
+        return tensor.reshape((np.prod(qid_shape, dtype=int).item(),) * 2)
 
     def _unitary_(self) -> Union[np.ndarray, NotImplementedType]:
         sub_matrix = protocols.unitary(self.sub_operation, None)

cirq/ops/dense_pauli_string.py

@@ -374,7 +374,7 @@ def frozen(self) -> 'DensePauliString':
     def copy(
         self,
         coefficient: Optional[complex] = None,
-        pauli_mask: Union[None, Iterable[int], np.ndarray] = None,
+        pauli_mask: Union[None, str, Iterable[int], np.ndarray] = None,
     ) -> 'DensePauliString':
         if pauli_mask is None and (coefficient is None or coefficient == self.coefficient):
             return self
@@ -446,7 +446,7 @@ def __imul__(self, other):
     def copy(
         self,
         coefficient: Optional[complex] = None,
-        pauli_mask: Union[None, Iterable[int], np.ndarray] = None,
+        pauli_mask: Union[None, str, Iterable[int], np.ndarray] = None,
     ) -> 'MutableDensePauliString':
         return MutableDensePauliString(
             coefficient=self.coefficient if coefficient is None else coefficient,
@@ -549,5 +549,5 @@ def _vectorized_pauli_mul_phase(
     t -= 1
 
     # Result is i raised to the sum of the per-term phase exponents.
-    s = int(np.sum(t, dtype=np.uint8) & 3)
+    s = int(np.sum(t, dtype=np.uint8).item() & 3)
     return 1j ** s

cirq/ops/diagonal_gate.py

@@ -18,7 +18,8 @@
 passed as a list.
 """
 
-from typing import AbstractSet, Any, Tuple, Iterator, List, Sequence, TYPE_CHECKING, Union
+from typing import AbstractSet, Any, Iterator, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+
 import numpy as np
 import sympy
 
@@ -30,7 +31,7 @@
     import cirq
 
 
-def _fast_walsh_hadamard_transform(a: Tuple[Any, ...]) -> np.array:
+def _fast_walsh_hadamard_transform(a: Tuple[Any, ...]) -> np.ndarray:
     """Fast Walsh–Hadamard Transform of an array."""
     h = 1
     a_ = np.array(a)
@@ -99,7 +100,7 @@ def _resolve_parameters_(
     def _has_unitary_(self) -> bool:
         return not self._is_parameterized_()
 
-    def _unitary_(self) -> np.ndarray:
+    def _unitary_(self) -> Optional[np.ndarray]:
         if self._is_parameterized_():
             return None
         return np.diag([np.exp(1j * angle) for angle in self._diag_angles_radians])

cirq/ops/eigen_gate.py

@@ -188,7 +188,7 @@ def _diagram_exponent(
 
         # Canonicalize the rounded exponent into (-period/2, period/2].
         if args.precision is not None:
-            result = np.around(result, args.precision)
+            result = np.around(result, args.precision).item()
         h = diagram_period / 2
         if not (-h < result <= h):
             result = h - result

cirq/ops/identity.py

@@ -73,7 +73,7 @@ def _has_unitary_(self) -> bool:
         return True
 
     def _unitary_(self) -> np.ndarray:
-        return np.identity(np.prod(self._qid_shape, dtype=int))
+        return np.identity(np.prod(self._qid_shape, dtype=int).item())
 
     def _apply_unitary_(self, args: 'protocols.ApplyUnitaryArgs') -> Optional[np.ndarray]:
         return args.target_tensor

cirq/ops/pauli_interaction_gate.py

@@ -112,14 +112,16 @@ def _decompose_(self, qubits: Sequence['cirq.Qid']) -> 'cirq.OP_TREE':
     def _circuit_diagram_info_(
         self, args: 'cirq.CircuitDiagramInfoArgs'
     ) -> 'cirq.CircuitDiagramInfo':
-        labels = cast(
-            Dict[pauli_gates.Pauli, np.ndarray],
-            {pauli_gates.X: 'X', pauli_gates.Y: 'Y', pauli_gates.Z: '@'},
-        )
+        labels: Dict['cirq.Pauli', str] = {
+            pauli_gates.X: 'X',
+            pauli_gates.Y: 'Y',
+            pauli_gates.Z: '@',
+        }
         l0 = labels[self.pauli0]
         l1 = labels[self.pauli1]
         # Add brackets around letter if inverted
-        l0, l1 = (f'(-{l})' if inv else l for l, inv in ((l0, self.invert0), (l1, self.invert1)))
+        l0 = f'(-{l0})' if self.invert0 else l0
+        l1 = f'(-{l1})' if self.invert1 else l1
         return protocols.CircuitDiagramInfo(
             wire_symbols=(l0, l1), exponent=self._diagram_exponent(args)
         )

cirq/ops/pauli_sum_exponential.py

@@ -104,7 +104,7 @@ def matrix(self) -> np.ndarray:
         """
         if protocols.is_parameterized(self._exponent):
             raise ValueError("Exponent should not parameterized.")
-        ret = 1
+        ret = np.ones(1)
         for pauli_string_exp in self:
             ret = np.kron(ret, protocols.unitary(pauli_string_exp))
         return ret

cirq/ops/phased_x_z_gate.py

@@ -108,7 +108,7 @@ def _value_equality_values_(self):
         )
 
     @staticmethod
-    def from_matrix(mat: np.array) -> 'cirq.PhasedXZGate':
+    def from_matrix(mat: np.ndarray) -> 'cirq.PhasedXZGate':
         pre_phase, rotation, post_phase = linalg.deconstruct_single_qubit_matrix_into_angles(mat)
         pre_phase /= np.pi
         post_phase /= np.pi

cirq/ops/raw_types.py

@@ -552,7 +552,7 @@ def _commutes_(
 
         m12 = protocols.unitary_protocol.unitary(circuit12, default=None)
         m21 = protocols.unitary_protocol.unitary(circuit21, default=None)
-        if m12 is None:
+        if m12 is None or m21 is None:
             return NotImplemented
 
         return np.allclose(m12, m21, atol=atol)

cirq/ops/two_qubit_diagonal_gate.py

@@ -64,7 +64,7 @@ def _resolve_parameters_(
     def _has_unitary_(self) -> bool:
         return not self._is_parameterized_()
 
-    def _unitary_(self) -> np.ndarray:
+    def _unitary_(self) -> Optional[np.ndarray]:
         if self._is_parameterized_():
             return None
         return np.diag([np.exp(1j * angle) for angle in self._diag_angles_radians])

cirq/qis/states.py

@@ -682,7 +682,7 @@ def validate_normalized_state_vector(
     state_vector: np.ndarray,
     *,  # Force keyword arguments
     qid_shape: Tuple[int, ...],
-    dtype: Optional[Type[np.number]] = None,
+    dtype: Optional[np.dtype] = None,
     atol: float = 1e-7,
 ) -> None:
     """Checks that the given state vector is valid.
@@ -754,7 +754,7 @@ def to_valid_density_matrix(
     num_qubits: Optional[int] = None,
     *,  # Force keyword arguments
     qid_shape: Optional[Tuple[int, ...]] = None,
-    dtype: Optional[Type[np.number]] = None,
+    dtype: Optional[np.dtype] = None,
     atol: float = 1e-7,
 ) -> np.ndarray:
     """Verifies the density_matrix_rep is valid and converts it to ndarray form.
@@ -888,9 +888,7 @@ def one_hot(
     return result
 
 
-def eye_tensor(
-    half_shape: Tuple[int, ...], *, dtype: Type[np.number]  # Force keyword args
-) -> np.array:
+def eye_tensor(half_shape: Tuple[int, ...], *, dtype: np.dtype) -> np.ndarray:
     """Returns an identity matrix reshaped into a tensor.
 
     Args:
@@ -902,6 +900,6 @@ def eye_tensor(
     Returns:
         The created numpy array with shape `half_shape + half_shape`.
     """
-    identity = np.eye(np.prod(half_shape, dtype=int), dtype=dtype)
+    identity = np.eye(np.prod(half_shape, dtype=int).item(), dtype=dtype)
     identity.shape = half_shape * 2
     return identity