scipy 1.15.3__cp312-cp312-win_amd64.whl → 1.16.0rc2__cp312-cp312-win_amd64.whl

scipy/_lib/array_api_compat/common/_fft.py

@@ -1,168 +1,195 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Union, Optional, Literal
+from collections.abc import Sequence
+from typing import Literal, TypeAlias
 
-if TYPE_CHECKING:
-    from ._typing import Device, ndarray
-    from collections.abc import Sequence
+from ._typing import Array, Device, DType, Namespace
+
+_Norm: TypeAlias = Literal["backward", "ortho", "forward"]
 
 # Note: NumPy fft functions improperly upcast float32 and complex64 to
 # complex128, which is why we require wrapping them all here.
 
 def fft(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.fft(x, n=n, axis=axis, norm=norm)
     if x.dtype in [xp.float32, xp.complex64]:
         return res.astype(xp.complex64)
     return res
 
 def ifft(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.ifft(x, n=n, axis=axis, norm=norm)
     if x.dtype in [xp.float32, xp.complex64]:
         return res.astype(xp.complex64)
     return res
 
 def fftn(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.fftn(x, s=s, axes=axes, norm=norm)
     if x.dtype in [xp.float32, xp.complex64]:
         return res.astype(xp.complex64)
     return res
 
 def ifftn(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.ifftn(x, s=s, axes=axes, norm=norm)
     if x.dtype in [xp.float32, xp.complex64]:
         return res.astype(xp.complex64)
     return res
 
 def rfft(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.rfft(x, n=n, axis=axis, norm=norm)
     if x.dtype == xp.float32:
         return res.astype(xp.complex64)
     return res
 
 def irfft(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.irfft(x, n=n, axis=axis, norm=norm)
     if x.dtype == xp.complex64:
         return res.astype(xp.float32)
     return res
 
 def rfftn(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.rfftn(x, s=s, axes=axes, norm=norm)
     if x.dtype == xp.float32:
         return res.astype(xp.complex64)
     return res
 
 def irfftn(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.irfftn(x, s=s, axes=axes, norm=norm)
     if x.dtype == xp.complex64:
         return res.astype(xp.float32)
     return res
 
 def hfft(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.hfft(x, n=n, axis=axis, norm=norm)
     if x.dtype in [xp.float32, xp.complex64]:
         return res.astype(xp.float32)
     return res
 
 def ihfft(
-    x: ndarray,
+    x: Array,
     /,
-    xp,
+    xp: Namespace,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
-    norm: Literal["backward", "ortho", "forward"] = "backward",
-) -> ndarray:
+    norm: _Norm = "backward",
+) -> Array:
     res = xp.fft.ihfft(x, n=n, axis=axis, norm=norm)
     if x.dtype in [xp.float32, xp.complex64]:
         return res.astype(xp.complex64)
     return res
 
-def fftfreq(n: int, /, xp, *, d: float = 1.0, device: Optional[Device] = None) -> ndarray:
+def fftfreq(
+    n: int,
+    /,
+    xp: Namespace,
+    *,
+    d: float = 1.0,
+    dtype: DType | None = None,
+    device: Device | None = None,
+) -> Array:
     if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
-    return xp.fft.fftfreq(n, d=d)
+    res = xp.fft.fftfreq(n, d=d)
+    if dtype is not None:
+        return res.astype(dtype)
+    return res
 
-def rfftfreq(n: int, /, xp, *, d: float = 1.0, device: Optional[Device] = None) -> ndarray:
+def rfftfreq(
+    n: int,
+    /,
+    xp: Namespace,
+    *,
+    d: float = 1.0,
+    dtype: DType | None = None,
+    device: Device | None = None,
+) -> Array:
     if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
-    return xp.fft.rfftfreq(n, d=d)
+    res = xp.fft.rfftfreq(n, d=d)
+    if dtype is not None:
+        return res.astype(dtype)
+    return res
 
-def fftshift(x: ndarray, /, xp, *, axes: Union[int, Sequence[int]] = None) -> ndarray:
+def fftshift(
+    x: Array, /, xp: Namespace, *, axes: int | Sequence[int] | None = None
+) -> Array:
     return xp.fft.fftshift(x, axes=axes)
 
-def ifftshift(x: ndarray, /, xp, *, axes: Union[int, Sequence[int]] = None) -> ndarray:
+def ifftshift(
+    x: Array, /, xp: Namespace, *, axes: int | Sequence[int] | None = None
+) -> Array:
     return xp.fft.ifftshift(x, axes=axes)
 
 __all__ = [
@@ -181,3 +208,6 @@ __all__ = [
     "fftshift",
     "ifftshift",
 ]
+
+def __dir__() -> list[str]:
+    return __all__