Use np.complexfloating for dtypes that should be complex (#5488)

This changes the type signature for dtypes that should only be `np.complex64` or `np.complex128`.

There are also dtypes in `qis.states.py` but I've not changed these over, as they are often called the `dtype` of a `np.array`, and so they will often need to be cast.

This fixes some `check/mypy --next` errors.

Technically a breaking change in type signature (more restrictive), but we yell when types are not complex.

Author: dabacon

cirq-core/cirq/circuits/circuit.py

@@ -64,7 +64,6 @@
 
 if TYPE_CHECKING:
     import cirq
-    from numpy.typing import DTypeLike
 
 
 _TGate = TypeVar('_TGate', bound='cirq.Gate')
@@ -999,7 +998,7 @@ def unitary(
         qubit_order: 'cirq.QubitOrderOrList' = ops.QubitOrder.DEFAULT,
         qubits_that_should_be_present: Iterable['cirq.Qid'] = (),
         ignore_terminal_measurements: bool = True,
-        dtype: 'DTypeLike' = np.complex64,
+        dtype: Type[np.complexfloating] = np.complex64,
     ) -> np.ndarray:
         """Converts the circuit into a unitary matrix, if possible.
 
@@ -1089,7 +1088,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['DTypeLike'] = None,
+        dtype: Optional[Type[np.complexfloating]] = None,
         param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
         seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     ) -> np.ndarray:
@@ -2656,7 +2655,7 @@ def _apply_unitary_circuit(
     circuit: 'cirq.AbstractCircuit',
     state: np.ndarray,
     qubits: Tuple['cirq.Qid', ...],
-    dtype: 'DTypeLike',
+    dtype: Type[np.complexfloating],
 ) -> 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, Union
+from typing import Any, Callable, List, Optional, Sequence, Tuple, Type, TYPE_CHECKING, Union
 
 import numpy as np
 
@@ -24,7 +24,6 @@
 
 if TYPE_CHECKING:
     import cirq
-    from numpy.typing import DTypeLike
 
 
 class _BufferedDensityMatrix(qis.QuantumStateRepresentation):
@@ -58,7 +57,7 @@ def create(
         *,
         initial_state: Union[np.ndarray, 'cirq.STATE_VECTOR_LIKE'] = 0,
         qid_shape: Optional[Tuple[int, ...]] = None,
-        dtype: Optional['DTypeLike'] = None,
+        dtype: Optional[Type[np.complexfloating]] = None,
         buffer: Optional[List[np.ndarray]] = None,
     ):
         """Creates a buffered density matrix with the requested state.
@@ -252,7 +251,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: 'DTypeLike' = np.complex64,
+        dtype: Type[np.complexfloating] = np.complex64,
         classical_data: Optional['cirq.ClassicalDataStore'] = None,
     ):
         """Inits DensityMatrixSimulationState.

cirq-core/cirq/sim/density_matrix_simulator.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Simulator for density matrices that simulates noisy quantum circuits."""
-from typing import Any, Dict, TYPE_CHECKING, Tuple, Union, Sequence, Optional, List
+from typing import Any, Dict, List, Optional, Sequence, Type, TYPE_CHECKING, Tuple, Union
 
 import numpy as np
 
@@ -22,7 +22,6 @@
 
 if TYPE_CHECKING:
     import cirq
-    from numpy.typing import DTypeLike
 
 
 class DensityMatrixSimulator(
@@ -116,7 +115,7 @@ class DensityMatrixSimulator(
     def __init__(
         self,
         *,
-        dtype: 'DTypeLike' = np.complex64,
+        dtype: Type[np.complexfloating] = np.complex64,
         noise: 'cirq.NOISE_MODEL_LIKE' = None,
         seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
         split_untangled_states: bool = True,
@@ -250,7 +249,7 @@ def __init__(
         self,
         sim_state: 'cirq.SimulationStateBase[cirq.DensityMatrixSimulationState]',
         simulator: 'cirq.DensityMatrixSimulator' = None,
-        dtype: 'DTypeLike' = np.complex64,
+        dtype: Type[np.complexfloating] = np.complex64,
     ):
         """DensityMatrixStepResult.
 

cirq-core/cirq/sim/density_matrix_simulator_test.py

@@ -17,7 +17,7 @@
 Filename is a reference to multiplexing.
 """
 
-from typing import cast, List, Optional, Sequence, TYPE_CHECKING, Union
+from typing import cast, List, Optional, Sequence, Type, TYPE_CHECKING, Union
 
 import numpy as np
 
@@ -53,7 +53,7 @@ def sample(
     noise: 'cirq.NOISE_MODEL_LIKE' = None,
     param_resolver: Optional['cirq.ParamResolver'] = None,
     repetitions: int = 1,
-    dtype: 'DTypeLike' = np.complex64,
+    dtype: Type[np.complexfloating] = np.complex64,
     seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
 ) -> 'cirq.Result':
     """Simulates sampling from the given circuit.
@@ -108,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: 'DTypeLike' = np.complex64,
+    dtype: Type[np.complexfloating] = np.complex64,
     seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
 ) -> 'np.ndarray':
     """Returns the state vector resulting from acting operations on a state.
@@ -178,7 +178,7 @@ def sample_sweep(
     *,
     noise: 'cirq.NOISE_MODEL_LIKE' = None,
     repetitions: int = 1,
-    dtype: 'DTypeLike' = np.complex64,
+    dtype: Type[np.complexfloating] = np.complex64,
     seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
 ) -> Sequence['cirq.Result']:
     """Runs the supplied Circuit, mimicking quantum hardware.
@@ -224,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: 'DTypeLike' = np.complex64,
+    dtype: Type[np.complexfloating] = np.complex64,
     seed: Optional[Union[int, np.random.RandomState]] = None,
     ignore_measurement_results: bool = True,
 ) -> 'np.ndarray':

cirq-core/cirq/sim/mux.py

@@ -26,6 +26,7 @@
     Optional,
     Sequence,
     Tuple,
+    Type,
     TypeVar,
     TYPE_CHECKING,
 )
@@ -93,7 +94,7 @@ class SimulatorBase(
     def __init__(
         self,
         *,
-        dtype: 'DTypeLike' = np.complex64,
+        dtype: Type[np.complexfloating] = 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/simulator_base.py

@@ -14,7 +14,7 @@
 
 """A simulator that uses numpy's einsum for sparse matrix operations."""
 
-from typing import Any, Iterator, List, TYPE_CHECKING, Union, Sequence, Optional
+from typing import Any, Iterator, List, TYPE_CHECKING, Union, Sequence, Type, Optional
 
 import numpy as np
 
@@ -24,7 +24,6 @@
 
 if TYPE_CHECKING:
     import cirq
-    from numpy.typing import DTypeLike
 
 
 class Simulator(
@@ -127,7 +126,7 @@ class Simulator(
     def __init__(
         self,
         *,
-        dtype: 'DTypeLike' = np.complex64,
+        dtype: Type[np.complexfloating] = np.complex64,
         noise: 'cirq.NOISE_MODEL_LIKE' = None,
         seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
         split_untangled_states: bool = True,
@@ -231,7 +230,7 @@ def __init__(
         self,
         sim_state: 'cirq.SimulationStateBase[cirq.StateVectorSimulationState]',
         simulator: 'cirq.Simulator' = None,
-        dtype: 'DTypeLike' = np.complex64,
+        dtype: Type[np.complexfloating] = np.complex64,
     ):
         """Results of a step of the simulator.