Clean up redundant complex type unions (#7041)

* Clean up redundant complex type unions

* format

* fix numpy ufunc, one mypy err

* Additional cleanup

Author: daxfohl

cirq-core/cirq/interop/quirk/cells/parse.py

@@ -86,7 +86,7 @@ def classify(e: str) -> Union[str, float]:
     return _merge_scientific_float_tokens(g for g in result if g.strip())
 
 
-_ResolvedToken = Union[sympy.Expr, int, float, complex]
+_ResolvedToken = Union[sympy.Expr, complex]
 
 
 class _CustomQuirkOperationToken:

cirq-core/cirq/linalg/combinators.py

@@ -25,7 +25,7 @@
     from numpy.typing import DTypeLike, ArrayLike
 
 
-def kron(*factors: Union[np.ndarray, complex, float], shape_len: int = 2) -> np.ndarray:
+def kron(*factors: Union[np.ndarray, complex], shape_len: int = 2) -> np.ndarray:
     """Computes the kronecker product of a sequence of values.
 
     A *args version of lambda args: functools.reduce(np.kron, args).
@@ -56,7 +56,7 @@ def kron(*factors: Union[np.ndarray, complex, float], shape_len: int = 2) -> np.
 )
 
 
-def kron_with_controls(*factors: Union[np.ndarray, complex, float]) -> np.ndarray:
+def kron_with_controls(*factors: Union[np.ndarray, complex]) -> np.ndarray:
     """Computes the kronecker product of a sequence of values and control tags.
 
     Use `cirq.CONTROL_TAG` to represent controls. Any entry of the output

cirq-core/cirq/linalg/decompositions.py

@@ -283,13 +283,7 @@ class AxisAngleDecomposition:
     rotation axis, and g is the global phase.
     """
 
-    def __init__(
-        self,
-        *,
-        angle: float,
-        axis: Tuple[float, float, float],
-        global_phase: Union[int, float, complex],
-    ):
+    def __init__(self, *, angle: float, axis: Tuple[float, float, float], global_phase: complex):
         if not np.isclose(np.linalg.norm(axis, 2), 1, atol=1e-8):
             raise ValueError('Axis vector must be normalized.')
         self.global_phase = complex(global_phase)
@@ -634,7 +628,7 @@ def scatter_plot_normalized_kak_interaction_coefficients(
         ax = cast(mplot3d.axes3d.Axes3D, fig.add_subplot(1, 1, 1, projection='3d'))
 
     def coord_transform(
-        pts: Union[List[Tuple[int, int, int]], np.ndarray]
+        pts: Union[List[Tuple[int, int, int]], np.ndarray],
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         if len(pts) == 0:
             return np.array([]), np.array([]), np.array([])

cirq-core/cirq/linalg/tolerance.py

@@ -33,7 +33,7 @@ def all_near_zero(a: 'ArrayLike', *, atol: float = 1e-8) -> bool:
 
 
 def all_near_zero_mod(
-    a: Union[float, complex, Iterable[float], np.ndarray], period: float, *, atol: float = 1e-8
+    a: Union[float, Iterable[float], np.ndarray], period: float, *, atol: float = 1e-8
 ) -> bool:
     """Checks if the tensor's elements are all near multiples of the period.
 

cirq-core/cirq/ops/dense_pauli_string.py

@@ -116,7 +116,7 @@ def pauli_mask(self) -> np.ndarray:
         return self._pauli_mask
 
     @property
-    def coefficient(self) -> Union[sympy.Expr, complex]:
+    def coefficient(self) -> 'cirq.TParamValComplex':
         """A complex coefficient or symbol."""
         return self._coefficient
 
@@ -359,7 +359,7 @@ def __str__(self) -> str:
             coef = '+'
         elif self.coefficient == -1:
             coef = '-'
-        elif isinstance(self.coefficient, (complex, sympy.Symbol)):
+        elif isinstance(self.coefficient, (numbers.Complex, sympy.Symbol)):
             coef = f'{self.coefficient}*'
         else:
             coef = f'({self.coefficient})*'
@@ -403,7 +403,7 @@ def mutable_copy(self) -> 'MutableDensePauliString':
     @abc.abstractmethod
     def copy(
         self,
-        coefficient: Optional[Union[sympy.Expr, int, float, complex]] = None,
+        coefficient: Optional['cirq.TParamValComplex'] = None,
         pauli_mask: Union[None, str, Iterable[int], np.ndarray] = None,
     ) -> Self:
         """Returns a copy with possibly modified contents.
@@ -459,7 +459,7 @@ def frozen(self) -> 'DensePauliString':
 
     def copy(
         self,
-        coefficient: Optional[Union[sympy.Expr, int, float, complex]] = None,
+        coefficient: Optional['cirq.TParamValComplex'] = 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):
@@ -559,7 +559,7 @@ def __imul__(self, other):
 
     def copy(
         self,
-        coefficient: Optional[Union[sympy.Expr, int, float, complex]] = None,
+        coefficient: Optional['cirq.TParamValComplex'] = None,
         pauli_mask: Union[None, str, Iterable[int], np.ndarray] = None,
     ) -> 'MutableDensePauliString':
         return MutableDensePauliString(

cirq-core/cirq/ops/eigen_gate.py

@@ -359,7 +359,7 @@ def _parameter_names_(self) -> AbstractSet[str]:
 
     def _resolve_parameters_(self, resolver: 'cirq.ParamResolver', recursive: bool) -> 'EigenGate':
         exponent = resolver.value_of(self._exponent, recursive)
-        if isinstance(exponent, (complex, numbers.Complex)):
+        if isinstance(exponent, numbers.Complex):
             if isinstance(exponent, numbers.Real):
                 exponent = float(exponent)
             else:

cirq-core/cirq/ops/identity.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """IdentityGate."""
 
+import numbers
 from types import NotImplementedType
 from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Sequence, Union
 
@@ -72,7 +73,7 @@ def num_qubits(self) -> int:
         return len(self._qid_shape)
 
     def __pow__(self, power: Any) -> Any:
-        if isinstance(power, (int, float, complex, sympy.Basic)):
+        if isinstance(power, (numbers.Complex, sympy.Basic)):
             return self
         return NotImplemented
 
@@ -126,7 +127,7 @@ def _json_dict_(self) -> Dict[str, Any]:
     def _mul_with_qubits(self, qubits: Tuple['cirq.Qid', ...], other):
         if isinstance(other, raw_types.Operation):
             return other
-        if isinstance(other, (complex, float, int)):
+        if isinstance(other, numbers.Complex):
             from cirq.ops.pauli_string import PauliString
 
             return PauliString(coefficient=other)

cirq-core/cirq/ops/linear_combinations.py

@@ -48,7 +48,7 @@
 
 UnitPauliStringT = FrozenSet[Tuple[raw_types.Qid, pauli_gates.Pauli]]
 PauliSumLike = Union[
-    int, float, complex, PauliString, 'PauliSum', pauli_string.SingleQubitPauliStringGateOperation
+    complex, PauliString, 'PauliSum', pauli_string.SingleQubitPauliStringGateOperation
 ]
 document(
     PauliSumLike,

cirq-core/cirq/ops/pauli_string.py

@@ -28,7 +28,6 @@
     Optional,
     overload,
     Sequence,
-    SupportsComplex,
     Tuple,
     TYPE_CHECKING,
     TypeVar,
@@ -271,9 +270,7 @@ def __mul__(self, other: 'cirq.Operation') -> 'cirq.PauliString[Union[TKey, cirq
         pass
 
     @overload
-    def __mul__(
-        self, other: Union[complex, int, float, numbers.Number]
-    ) -> 'cirq.PauliString[TKey]':
+    def __mul__(self, other: complex) -> 'cirq.PauliString[TKey]':
         pass
 
     def __mul__(self, other):
@@ -308,10 +305,9 @@ def gate(self) -> 'cirq.DensePauliString':
         )
 
     def __rmul__(self, other) -> 'PauliString':
-        if isinstance(other, numbers.Number):
+        if isinstance(other, numbers.Complex):
             return PauliString(
-                qubit_pauli_map=self._qubit_pauli_map,
-                coefficient=self._coefficient * complex(cast(SupportsComplex, other)),
+                qubit_pauli_map=self._qubit_pauli_map, coefficient=self._coefficient * other
             )
 
         if isinstance(other, raw_types.Operation) and isinstance(other.gate, identity.IdentityGate):
@@ -321,10 +317,9 @@ def __rmul__(self, other) -> 'PauliString':
         return NotImplemented
 
     def __truediv__(self, other):
-        if isinstance(other, numbers.Number):
+        if isinstance(other, numbers.Complex):
             return PauliString(
-                qubit_pauli_map=self._qubit_pauli_map,
-                coefficient=self._coefficient / complex(cast(SupportsComplex, other)),
+                qubit_pauli_map=self._qubit_pauli_map, coefficient=self._coefficient / other
             )
         return NotImplemented
 
@@ -518,7 +513,7 @@ def _unitary_(self) -> Optional[np.ndarray]:
     def _apply_unitary_(self, args: 'protocols.ApplyUnitaryArgs'):
         if not self._has_unitary_():
             return None
-        assert isinstance(self.coefficient, complex)
+        assert isinstance(self.coefficient, numbers.Complex)
         if self.coefficient != 1:
             args.target_tensor *= self.coefficient
         return protocols.apply_unitaries([self[q].on(q) for q in self.qubits], self.qubits, args)
@@ -792,9 +787,11 @@ def __pos__(self) -> 'PauliString':
         return self
 
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
-        """Override behavior of numpy's exp method."""
+        """Override numpy behavior."""
         if ufunc == np.exp and len(inputs) == 1 and inputs[0] is self:
             return math.e**self
+        if ufunc == np.multiply and len(inputs) == 2 and inputs[1] is self:
+            return self * inputs[0]
         return NotImplemented
 
     def __pow__(self, power):
@@ -1174,14 +1171,14 @@ def _as_pauli_string(self) -> PauliString:
     def __mul__(self, other):
         if isinstance(other, SingleQubitPauliStringGateOperation):
             return self._as_pauli_string() * other._as_pauli_string()
-        if isinstance(other, (PauliString, complex, float, int)):
+        if isinstance(other, (PauliString, numbers.Complex)):
             return self._as_pauli_string() * other
         if (as_pauli_string := _try_interpret_as_pauli_string(other)) is not None:
             return self * as_pauli_string
         return NotImplemented
 
     def __rmul__(self, other):
-        if isinstance(other, (PauliString, complex, float, int)):
+        if isinstance(other, (PauliString, numbers.Complex)):
             return other * self._as_pauli_string()
         if (as_pauli_string := _try_interpret_as_pauli_string(other)) is not None:
             return as_pauli_string * self
@@ -1430,8 +1427,8 @@ def _imul_helper(self, other: 'cirq.PAULI_STRING_LIKE', sign: int):
                 pauli_int = _pauli_like_to_pauli_int(qubit, pauli_gate_like)
                 phase_log_i += self._imul_atom_helper(cast(TKey, qubit), pauli_int, sign)
             self.coefficient *= 1j ** (phase_log_i & 3)
-        elif isinstance(other, numbers.Number):
-            self.coefficient *= complex(cast(SupportsComplex, other))
+        elif isinstance(other, numbers.Complex):
+            self.coefficient *= other
         elif isinstance(other, raw_types.Operation) and isinstance(
             other.gate, identity.IdentityGate
         ):

cirq-core/cirq/ops/pauli_string_phasor.py

@@ -392,14 +392,14 @@ def _resolve_parameters_(
     ) -> 'PauliStringPhasorGate':
         exponent_neg = resolver.value_of(self.exponent_neg, recursive)
         exponent_pos = resolver.value_of(self.exponent_pos, recursive)
-        if isinstance(exponent_neg, (complex, numbers.Complex)):
+        if isinstance(exponent_neg, numbers.Complex):
             if isinstance(exponent_neg, numbers.Real):
                 exponent_neg = float(exponent_neg)
             else:
                 raise ValueError(
                     f'PauliStringPhasorGate does not support complex exponent {exponent_neg}'
                 )
-        if isinstance(exponent_pos, (complex, numbers.Complex)):
+        if isinstance(exponent_pos, numbers.Complex):
             if isinstance(exponent_pos, numbers.Real):
                 exponent_pos = float(exponent_pos)
             else: