Fix some mypy --next numpy typing errors

cirq-core/cirq/circuits/circuit.py

@@ -64,6 +64,7 @@
 
 if TYPE_CHECKING:
     import cirq
+    from numpy.typing import DTypeLike
 
 
 _TGate = TypeVar('_TGate', bound='cirq.Gate')
@@ -998,7 +999,7 @@ def unitary(
         qubit_order: 'cirq.QubitOrderOrList' = ops.QubitOrder.DEFAULT,
         qubits_that_should_be_present: Iterable['cirq.Qid'] = (),
         ignore_terminal_measurements: bool = True,
-        dtype: Type[np.number] = np.complex64,
+        dtype: 'DTypeLike' = np.complex64,
     ) -> np.ndarray:
         """Converts the circuit into a unitary matrix, if possible.
 
@@ -1088,7 +1089,7 @@ def final_state_vector(
         qubit_order: 'cirq.QubitOrderOrList' = ops.QubitOrder.DEFAULT,
         qubits_that_should_be_present: Iterable['cirq.Qid'] = (),
         ignore_terminal_measurements: Optional[bool] = None,
-        dtype: Optional[Type[np.number]] = None,
+        dtype: Optional['DTypeLike'] = None,
         param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
         seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     ) -> np.ndarray:
@@ -2589,7 +2590,7 @@ def _apply_unitary_circuit(
     circuit: 'cirq.AbstractCircuit',
     state: np.ndarray,
     qubits: Tuple['cirq.Qid', ...],
-    dtype: Type[np.number],
+    dtype: 'DTypeLike',
 ) -> np.ndarray:
     """Applies a circuit's unitary effect to the given vector or matrix.
 

cirq-core/cirq/sim/density_matrix_simulation_state.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Objects and methods for acting efficiently on a density matrix."""
 
-from typing import Any, Callable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Type, Union
+from typing import Any, Callable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 
@@ -252,7 +252,7 @@ def __init__(
         prng: Optional[np.random.RandomState] = None,
         qubits: Optional[Sequence['cirq.Qid']] = None,
         initial_state: Union[np.ndarray, 'cirq.STATE_VECTOR_LIKE'] = 0,
-        dtype: Type[np.number] = np.complex64,
+        dtype: 'DTypeLike' = np.complex64,
         classical_data: Optional['cirq.ClassicalDataStore'] = None,
     ):
         """Inits DensityMatrixSimulationState.

cirq-core/cirq/sim/density_matrix_utils.py

@@ -94,7 +94,7 @@ def measure_density_matrix(
     density_matrix: np.ndarray,
     indices: Sequence[int],
     qid_shape: Optional[Tuple[int, ...]] = None,
-    out: np.ndarray = None,
+    out: Optional[np.ndarray] = None,
     seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
 ) -> Tuple[List[int], np.ndarray]:
     """Performs a measurement of the density matrix in the computational basis.
@@ -142,13 +142,16 @@ def measure_density_matrix(
         num_qubits = len(qid_shape)
     meas_shape = _indices_shape(qid_shape, indices)
 
+    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]
+    )
+
     if len(indices) == 0:
-        if out is None:
-            out = np.copy(density_matrix)
-        elif out is not density_matrix:
-            np.copyto(dst=out, src=density_matrix)
-        return ([], out)
-        # Final else: if out is matrix then matrix will be modified in place.
+        return ([], arrout)
 
     prng = value.parse_random_state(seed)
 
@@ -169,21 +172,15 @@ def measure_density_matrix(
     # Remove ellipses from last element of
     mask[result_slice * 2] = False
 
-    if out is None:
-        out = np.copy(density_matrix)
-    elif out is not density_matrix:
-        np.copyto(dst=out, src=density_matrix)
-    # Final else: if out is matrix then matrix will be modified in place.
-
     # Potentially reshape to tensor, and then set masked values to 0.
-    out.shape = qid_shape * 2
-    out[mask] = 0
+    arrout.shape = qid_shape * 2
+    arrout[mask] = 0
 
     # Restore original shape (if necessary) and renormalize.
-    out.shape = initial_shape
-    out /= probs[result]
+    arrout.shape = initial_shape
+    arrout /= probs[result]
 
-    return measurement_bits, out
+    return measurement_bits, arrout
 
 
 def _probs(

cirq-core/cirq/sim/mux.py

@@ -17,7 +17,7 @@
 Filename is a reference to multiplexing.
 """
 
-from typing import cast, List, Optional, Sequence, Type, TYPE_CHECKING, Union
+from typing import cast, List, Optional, Sequence, TYPE_CHECKING, Union
 
 import numpy as np
 
@@ -29,6 +29,7 @@
 
 if TYPE_CHECKING:
     import cirq
+    from numpy.typing import DTypeLike
 
 CIRCUIT_LIKE = Union[circuits.Circuit, ops.Gate, ops.OP_TREE]
 document(
@@ -52,7 +53,7 @@ def sample(
     noise: 'cirq.NOISE_MODEL_LIKE' = None,
     param_resolver: Optional['cirq.ParamResolver'] = None,
     repetitions: int = 1,
-    dtype: Type[np.number] = np.complex64,
+    dtype: 'DTypeLike' = np.complex64,
     seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
 ) -> 'cirq.Result':
     """Simulates sampling from the given circuit.
@@ -107,7 +108,7 @@ def final_state_vector(
     param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
     qubit_order: 'cirq.QubitOrderOrList' = ops.QubitOrder.DEFAULT,
     ignore_terminal_measurements: bool = False,
-    dtype: Type[np.number] = np.complex64,
+    dtype: 'DTypeLike' = np.complex64,
     seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
 ) -> 'np.ndarray':
     """Returns the state vector resulting from acting operations on a state.
@@ -177,7 +178,7 @@ def sample_sweep(
     *,
     noise: 'cirq.NOISE_MODEL_LIKE' = None,
     repetitions: int = 1,
-    dtype: Type[np.number] = np.complex64,
+    dtype: 'DTypeLike' = np.complex64,
     seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
 ) -> Sequence['cirq.Result']:
     """Runs the supplied Circuit, mimicking quantum hardware.
@@ -223,7 +224,7 @@ def final_density_matrix(
     initial_state: 'cirq.STATE_VECTOR_LIKE' = 0,
     param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
     qubit_order: 'cirq.QubitOrderOrList' = ops.QubitOrder.DEFAULT,
-    dtype: Type[np.number] = np.complex64,
+    dtype: 'DTypeLike' = np.complex64,
     seed: Optional[Union[int, np.random.RandomState]] = None,
     ignore_measurement_results: bool = True,
 ) -> 'np.ndarray':

cirq-core/cirq/sim/simulator_base.py

@@ -26,7 +26,6 @@
     Optional,
     Sequence,
     Tuple,
-    Type,
     TypeVar,
     TYPE_CHECKING,
 )
@@ -50,6 +49,7 @@
 
 if TYPE_CHECKING:
     import cirq
+    from numpy.typing import DTypeLike
 
 
 TStepResultBase = TypeVar('TStepResultBase', bound='StepResultBase')
@@ -93,7 +93,7 @@ class SimulatorBase(
     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,
         split_untangled_states: bool = False,

cirq-core/cirq/sim/sparse_simulator.py

@@ -14,7 +14,7 @@
 
 """A simulator that uses numpy's einsum for sparse matrix operations."""
 
-from typing import Any, Iterator, List, Type, TYPE_CHECKING, Union, Sequence, Optional
+from typing import Any, Iterator, List, TYPE_CHECKING, Union, Sequence, Optional
 
 import numpy as np
 
@@ -127,7 +127,7 @@ class Simulator(
     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,
         split_untangled_states: bool = True,