Fix most numpy errors in cirq/qis and friends (#4002)

Part of #3767

Author: mpharrigan

cirq/interop/quirk/cells/parse.py

@@ -70,7 +70,7 @@ def _segment_by(seq: Iterable[T], *, key: Callable[[T], Any]) -> Iterator[List[T
 
 
 def _tokenize(text: str) -> List[str]:
-    def classify(e: str) -> str:
+    def classify(e: str) -> Union[str, float]:
         assert e.strip() != ''  # Because _segment_by drops empty entries.
         if re.match(r'[.0-9]', e):
             return "#"

cirq/qis/measures.py

@@ -138,9 +138,9 @@ def fidelity(
         )
     state1 = quantum_state(state1, qid_shape=qid_shape, validate=validate, atol=atol)
     state2 = quantum_state(state2, qid_shape=qid_shape, validate=validate, atol=atol)
-    state1 = state1.density_matrix() if state1._is_density_matrix() else state1.state_vector()
-    state2 = state2.density_matrix() if state2._is_density_matrix() else state2.state_vector()
-    return _fidelity_state_vectors_or_density_matrices(state1, state2)
+    state1_arr = state1.state_vector_or_density_matrix()
+    state2_arr = state2.state_vector_or_density_matrix()
+    return _fidelity_state_vectors_or_density_matrices(state1_arr, state2_arr)
 
 
 def _numpy_arrays_to_state_vectors_or_density_matrices(
@@ -152,10 +152,10 @@ def _numpy_arrays_to_state_vectors_or_density_matrices(
 ) -> Tuple[np.ndarray, np.ndarray]:
     if state1.ndim > 2 or (state1.ndim == 2 and state1.shape[0] != state1.shape[1]):
         # State tensor, convert to state vector
-        state1 = np.reshape(state1, (np.prod(state1.shape),))
+        state1 = np.reshape(state1, (np.prod(state1.shape).item(),))
     if state2.ndim > 2 or (state2.ndim == 2 and state2.shape[0] != state2.shape[1]):
         # State tensor, convert to state vector
-        state2 = np.reshape(state2, (np.prod(state2.shape),))
+        state2 = np.reshape(state2, (np.prod(state2.shape).item(),))
     if state1.ndim == 2 and state2.ndim == 2:
         # Must be square matrices
         if state1.shape == state2.shape:
@@ -168,26 +168,26 @@ def _numpy_arrays_to_state_vectors_or_density_matrices(
                 )
             if state1.shape == qid_shape:
                 # State tensors, convert to state vectors
-                state1 = np.reshape(state1, (np.prod(qid_shape),))
-                state2 = np.reshape(state2, (np.prod(qid_shape),))
+                state1 = np.reshape(state1, (np.prod(qid_shape).item(),))
+                state2 = np.reshape(state2, (np.prod(qid_shape).item(),))
         elif state1.shape[0] < state2.shape[0]:
             # state1 is state tensor and state2 is density matrix.
             # Convert state1 to state vector
-            state1 = np.reshape(state1, (np.prod(state1.shape),))
+            state1 = np.reshape(state1, (np.prod(state1.shape).item(),))
         else:  # state1.shape[0] > state2.shape[0]
             # state2 is state tensor and state1 is density matrix.
             # Convert state2 to state vector
-            state2 = np.reshape(state2, (np.prod(state2.shape),))
+            state2 = np.reshape(state2, (np.prod(state2.shape).item(),))
     elif state1.ndim == 2 and state2.ndim < 2 and np.prod(state1.shape) == np.prod(state2.shape):
         # state1 is state tensor, convert to state vector
-        state1 = np.reshape(state1, (np.prod(state1.shape),))
+        state1 = np.reshape(state1, (np.prod(state1.shape).item(),))
     elif state1.ndim < 2 and state2.ndim == 2 and np.prod(state1.shape) == np.prod(state2.shape):
         # state2 is state tensor, convert to state vector
-        state2 = np.reshape(state2, (np.prod(state2.shape),))
+        state2 = np.reshape(state2, (np.prod(state2.shape).item(),))
 
     if validate:
-        dim1 = state1.shape[0] if state1.ndim == 2 else np.prod(state1.shape)
-        dim2 = state2.shape[0] if state2.ndim == 2 else np.prod(state2.shape)
+        dim1: int = state1.shape[0] if state1.ndim == 2 else np.prod(state1.shape).item()
+        dim2: int = state2.shape[0] if state2.ndim == 2 else np.prod(state2.shape).item()
         if dim1 != dim2:
             raise ValueError('Mismatched dimensions in given states: ' f'{dim1} and {dim2}.')
         if qid_shape is None:

cirq/qis/states.py

@@ -14,7 +14,7 @@
 """Classes and methods for quantum states."""
 
 import itertools
-from typing import Any, cast, Iterable, Optional, Sequence, TYPE_CHECKING, Tuple, Type, Union
+from typing import Any, cast, Iterable, Optional, Sequence, TYPE_CHECKING, Tuple, Union
 
 import numpy as np
 
@@ -23,6 +23,7 @@
 
 if TYPE_CHECKING:
     import cirq
+    from numpy.typing import DTypeLike
 
 DEFAULT_COMPLEX_DTYPE = np.complex64
 
@@ -61,7 +62,7 @@ def __init__(
         data: np.ndarray,
         qid_shape: Optional[Tuple[int, ...]] = None,
         *,  # Force keyword arguments
-        dtype: Optional[Type[np.number]] = None,
+        dtype: Optional['DTypeLike'] = None,
         validate: bool = True,
         atol: float = 1e-7,
     ) -> None:
@@ -87,7 +88,7 @@ def __init__(
             )
         self._data = data
         self._qid_shape = qid_shape
-        self._dim = np.prod(self.qid_shape, dtype=int)
+        self._dim = np.prod(self.qid_shape, dtype=int).item()
         if validate:
             self.validate(dtype=dtype, atol=atol)
 
@@ -102,7 +103,7 @@ def qid_shape(self) -> Tuple[int, ...]:
         return self._qid_shape
 
     @property
-    def dtype(self) -> np.ndarray:
+    def dtype(self) -> np.dtype:
         """The data type of the quantum state."""
         return self._data.dtype
 
@@ -136,15 +137,27 @@ def density_matrix(self) -> np.ndarray:
         """
         if not self._is_density_matrix():
             state_vector = self.state_vector()
+            assert state_vector is not None, 'only None if _is_density_matrix'
             return np.outer(state_vector, np.conj(state_vector))
         return self.data
 
+    def state_vector_or_density_matrix(self) -> np.ndarray:
+        """Return the state vector or density matrix of this state.
+
+        If the state is a denity matrix, return the density matrix. Otherwise, return the state
+        vector.
+        """
+        state_vector = self.state_vector()
+        if state_vector is not None:
+            return state_vector
+        return self.data
+
     def _is_density_matrix(self) -> bool:
         """Whether this quantum state is a density matrix."""
         return self.data.shape == (self._dim, self._dim)
 
     def validate(
-        self, *, dtype: Optional[Type[np.number]] = None, atol=1e-7  # Force keyword arguments
+        self, *, dtype: Optional['DTypeLike'] = None, atol=1e-7  # Force keyword arguments
     ) -> None:
         """Check if this quantum state is valid.
 
@@ -158,8 +171,13 @@ def validate(
         is_state_vector = self.data.shape == (self._dim,)
         is_state_tensor = self.data.shape == self.qid_shape
         if is_state_vector or is_state_tensor:
+            state_vector = self.state_vector()
+            assert state_vector is not None
             validate_normalized_state_vector(
-                self.state_vector(), qid_shape=self.qid_shape, dtype=dtype, atol=atol
+                state_vector,
+                qid_shape=self.qid_shape,
+                dtype=dtype,
+                atol=atol,
             )
         elif self._is_density_matrix():
             validate_density_matrix(
@@ -179,7 +197,7 @@ def quantum_state(
     *,  # Force keyword arguments
     copy: bool = False,
     validate: bool = True,
-    dtype: Optional[Type[np.number]] = None,
+    dtype: Optional['DTypeLike'] = None,
     atol: float = 1e-7,
 ) -> QuantumState:
     """Create a QuantumState object from a state-like object.
@@ -239,7 +257,7 @@ def quantum_state(
                 'Please specify the qid shape explicitly using '
                 'the qid_shape argument.'
             )
-        dim = np.prod(qid_shape, dtype=int)
+        dim = np.prod(qid_shape, dtype=int).item()
         if dtype is None:
             dtype = DEFAULT_COMPLEX_DTYPE
         data = one_hot(index=state, shape=(dim,), dtype=dtype)
@@ -279,7 +297,7 @@ def density_matrix(
     *,  # Force keyword arguments
     copy: bool = False,
     validate: bool = True,
-    dtype: Optional[Type[np.number]] = None,
+    dtype: Optional['DTypeLike'] = None,
     atol: float = 1e-7,
 ) -> QuantumState:
     """Create a QuantumState object from a density matrix.
@@ -494,7 +512,7 @@ def to_valid_state_vector(
     num_qubits: Optional[int] = None,
     *,  # Force keyword arguments
     qid_shape: Optional[Sequence[int]] = None,
-    dtype: Optional[Type[np.number]] = None,
+    dtype: Optional['DTypeLike'] = None,
     atol: float = 1e-7,
 ) -> np.ndarray:
     """Verifies the state_rep is valid and converts it to ndarray form.
@@ -555,7 +573,7 @@ def _state_like_to_state_tensor(
     *,
     state_like: 'cirq.STATE_VECTOR_LIKE',
     qid_shape: Tuple[int, ...],
-    dtype: Optional[Type[np.number]],
+    dtype: Optional['DTypeLike'],
     atol: float,
 ) -> np.ndarray:
 
@@ -619,7 +637,7 @@ def _amplitudes_to_validated_state_tensor(
     *,
     state_vector: np.ndarray,
     qid_shape: Tuple[int, ...],
-    dtype: Optional[Type[np.number]],
+    dtype: Optional['DTypeLike'],
     atol: float,
 ) -> np.ndarray:
     if dtype is None:
@@ -630,7 +648,7 @@ def _amplitudes_to_validated_state_tensor(
 
 
 def _qudit_values_to_state_tensor(
-    *, state_vector: np.ndarray, qid_shape: Tuple[int, ...], dtype: Optional[Type[np.number]]
+    *, state_vector: np.ndarray, qid_shape: Tuple[int, ...], dtype: Optional['DTypeLike']
 ) -> np.ndarray:
 
     for i in range(len(qid_shape)):
@@ -661,9 +679,9 @@ def _qudit_values_to_state_tensor(
 
 
 def _computational_basis_state_to_state_tensor(
-    *, state_rep: int, qid_shape: Tuple[int, ...], dtype: Optional[Type[np.number]]
+    *, state_rep: int, qid_shape: Tuple[int, ...], dtype: Optional['DTypeLike']
 ) -> np.ndarray:
-    n = np.prod(qid_shape, dtype=int)
+    n = np.prod(qid_shape, dtype=int).item()
     if not 0 <= state_rep < n:
         raise ValueError(
             f'Computational basis state is out of range.\n'
@@ -682,7 +700,7 @@ def validate_normalized_state_vector(
     state_vector: np.ndarray,
     *,  # Force keyword arguments
     qid_shape: Tuple[int, ...],
-    dtype: Optional[np.dtype] = None,
+    dtype: Optional['DTypeLike'] = None,
     atol: float = 1e-7,
 ) -> None:
     """Checks that the given state vector is valid.
@@ -754,7 +772,7 @@ def to_valid_density_matrix(
     num_qubits: Optional[int] = None,
     *,  # Force keyword arguments
     qid_shape: Optional[Tuple[int, ...]] = None,
-    dtype: Optional[np.dtype] = None,
+    dtype: Optional['DTypeLike'] = None,
     atol: float = 1e-7,
 ) -> np.ndarray:
     """Verifies the density_matrix_rep is valid and converts it to ndarray form.
@@ -799,7 +817,7 @@ def validate_density_matrix(
     density_matrix: np.ndarray,
     *,  # Force keyword arguments
     qid_shape: Tuple[int, ...],
-    dtype: Optional[Type[np.number]] = None,
+    dtype: Optional['DTypeLike'] = None,
     atol: float = 1e-7,
 ) -> None:
     """Checks that the given density matrix is valid.
@@ -867,7 +885,7 @@ def one_hot(
     index: Union[None, int, Sequence[int]] = None,
     shape: Union[int, Sequence[int]],
     value: Any = 1,
-    dtype: Type[np.number],
+    dtype: 'DTypeLike',
 ) -> np.ndarray:
     """Returns a numpy array with all 0s and a single non-zero entry(default 1).
 
@@ -888,7 +906,7 @@ def one_hot(
     return result
 
 
-def eye_tensor(half_shape: Tuple[int, ...], *, dtype: np.dtype) -> np.ndarray:
+def eye_tensor(half_shape: Tuple[int, ...], *, dtype: 'DTypeLike') -> np.ndarray:
     """Returns an identity matrix reshaped into a tensor.
 
     Args:

cirq/qis/states_test.py

@@ -50,6 +50,7 @@ def test_quantum_state():
     np.testing.assert_array_equal(state.state_vector(), state_vector_1)
     np.testing.assert_array_equal(state.state_tensor(), state_tensor_1)
     np.testing.assert_array_equal(state.density_matrix(), density_matrix_1)
+    np.testing.assert_array_equal(state.state_vector_or_density_matrix(), state_vector_1)
 
     state = cirq.QuantumState(state_tensor_1, qid_shape=(2, 2))
     assert state.data is state_tensor_1
@@ -58,6 +59,7 @@ def test_quantum_state():
     np.testing.assert_array_equal(state.state_vector(), state_vector_1)
     np.testing.assert_array_equal(state.state_tensor(), state_tensor_1)
     np.testing.assert_array_equal(state.density_matrix(), density_matrix_1)
+    np.testing.assert_array_equal(state.state_vector_or_density_matrix(), state_vector_1)
 
     state = cirq.QuantumState(density_matrix_1, qid_shape=(2, 2))
     assert state.data is density_matrix_1
@@ -66,6 +68,7 @@ def test_quantum_state():
     assert state.state_vector() is None
     assert state.state_tensor() is None
     np.testing.assert_array_equal(state.density_matrix(), density_matrix_1)
+    np.testing.assert_array_equal(state.state_vector_or_density_matrix(), density_matrix_1)
 
 
 def test_quantum_state_quantum_state():

cirq/sim/density_matrix_simulator.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Simulator for density matrices that simulates noisy quantum circuits."""
 import collections
-from typing import Any, Dict, Iterator, List, TYPE_CHECKING, Tuple, Type, Union
+from typing import Any, Dict, Iterator, List, TYPE_CHECKING, Tuple, Union
 
 import numpy as np
 
@@ -24,6 +24,7 @@
 
 if TYPE_CHECKING:
     import cirq
+    from numpy.typing import DTypeLike
 
 
 class DensityMatrixSimulator(
@@ -114,7 +115,7 @@ class DensityMatrixSimulator(
     def __init__(
         self,
         *,
-        dtype: Type[np.number] = np.complex64,
+        dtype: 'DTypeLike' = np.complex64,
         noise: 'cirq.NOISE_MODEL_LIKE' = None,
         seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
         ignore_measurement_results: bool = False,
@@ -320,7 +321,7 @@ def __init__(
         density_matrix: np.ndarray,
         measurements: Dict[str, np.ndarray],
         qubit_map: Dict[ops.Qid, int],
-        dtype: Type[np.number] = np.complex64,
+        dtype: 'DTypeLike' = np.complex64,
     ):
         """DensityMatrixStepResult.
 

cirq/study/result.py

@@ -356,8 +356,8 @@ def _pack_digits(digits: np.ndarray, pack_bits: str = 'auto') -> Tuple[str, bool
         # Do error checking here, otherwise the following logic will work
         # for both "auto" and "never".
 
-    if pack_bits == 'auto' and np.array_equal(digits, digits.astype(np.bool)):
-        return _pack_bits(digits.astype(np.bool)), True
+    if pack_bits == 'auto' and np.array_equal(digits, digits.astype(np.bool_)):
+        return _pack_bits(digits.astype(np.bool_)), True
 
     buffer = io.BytesIO()
     np.save(buffer, digits, allow_pickle=False)
@@ -404,4 +404,4 @@ def _unpack_digits(
 def _unpack_bits(packed_bits: str, dtype: str, shape: Sequence[int]) -> np.ndarray:
     bits_bytes = bytes.fromhex(packed_bits)
     bits = np.unpackbits(np.frombuffer(bits_bytes, dtype=np.uint8))
-    return bits[: np.prod(shape)].reshape(shape).astype(dtype)
+    return bits[: np.prod(shape).item()].reshape(shape).astype(dtype)

cirq/testing/consistent_qasm.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import warnings
-from typing import Any, List, Sequence
+from typing import Any, List, Sequence, Optional
 
 import numpy as np
 
@@ -81,6 +81,8 @@ def assert_qasm_is_consistent_with_unitary(val: Any):
             qasm_unitary, unitary, rtol=1e-8, atol=1e-8
         )
     except Exception as ex:
+        p_unitary: Optional[np.ndarray]
+        p_qasm_unitary: Optional[np.ndarray]
         if qasm_unitary is not None:
             p_unitary, p_qasm_unitary = linalg.match_global_phase(unitary, qasm_unitary)
         else:

cirq/value/product_state.py

@@ -131,7 +131,7 @@ def state_vector(self, qubit_order: 'cirq.QubitOrder' = None) -> np.ndarray:
         qubit_order = ops.QubitOrder.as_qubit_order(qubit_order)
         qubits = qubit_order.order_for(self.qubits)
 
-        mat = 1.0 + 0.0j
+        mat = np.ones(1, dtype=np.complex128)
         for qubit in qubits:
             oneq_state = self[qubit]
             state_vector = oneq_state.state_vector()
@@ -151,7 +151,7 @@ def projector(self, qubit_order: 'cirq.QubitOrder' = None) -> np.ndarray:
         qubit_order = ops.QubitOrder.as_qubit_order(qubit_order)
         qubits = qubit_order.order_for(self.qubits)
 
-        mat = 1.0 + 0.0j
+        mat = np.ones(1, dtype=np.complex128)
         for qubit in qubits:
             oneq_state = self[qubit]
             oneq_proj = oneq_state.projector()