_pocketfft.pyi 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. from collections.abc import Sequence
  2. from typing import Literal as L
  3. from typing import TypeAlias
  4. from numpy import complex128, float64
  5. from numpy._typing import ArrayLike, NDArray, _ArrayLikeNumber_co
  6. __all__ = [
  7. "fft",
  8. "ifft",
  9. "rfft",
  10. "irfft",
  11. "hfft",
  12. "ihfft",
  13. "rfftn",
  14. "irfftn",
  15. "rfft2",
  16. "irfft2",
  17. "fft2",
  18. "ifft2",
  19. "fftn",
  20. "ifftn",
  21. ]
  22. _NormKind: TypeAlias = L["backward", "ortho", "forward"] | None
  23. def fft(
  24. a: ArrayLike,
  25. n: int | None = ...,
  26. axis: int = ...,
  27. norm: _NormKind = ...,
  28. out: NDArray[complex128] | None = ...,
  29. ) -> NDArray[complex128]: ...
  30. def ifft(
  31. a: ArrayLike,
  32. n: int | None = ...,
  33. axis: int = ...,
  34. norm: _NormKind = ...,
  35. out: NDArray[complex128] | None = ...,
  36. ) -> NDArray[complex128]: ...
  37. def rfft(
  38. a: ArrayLike,
  39. n: int | None = ...,
  40. axis: int = ...,
  41. norm: _NormKind = ...,
  42. out: NDArray[complex128] | None = ...,
  43. ) -> NDArray[complex128]: ...
  44. def irfft(
  45. a: ArrayLike,
  46. n: int | None = ...,
  47. axis: int = ...,
  48. norm: _NormKind = ...,
  49. out: NDArray[float64] | None = ...,
  50. ) -> NDArray[float64]: ...
  51. # Input array must be compatible with `np.conjugate`
  52. def hfft(
  53. a: _ArrayLikeNumber_co,
  54. n: int | None = ...,
  55. axis: int = ...,
  56. norm: _NormKind = ...,
  57. out: NDArray[float64] | None = ...,
  58. ) -> NDArray[float64]: ...
  59. def ihfft(
  60. a: ArrayLike,
  61. n: int | None = ...,
  62. axis: int = ...,
  63. norm: _NormKind = ...,
  64. out: NDArray[complex128] | None = ...,
  65. ) -> NDArray[complex128]: ...
  66. def fftn(
  67. a: ArrayLike,
  68. s: Sequence[int] | None = ...,
  69. axes: Sequence[int] | None = ...,
  70. norm: _NormKind = ...,
  71. out: NDArray[complex128] | None = ...,
  72. ) -> NDArray[complex128]: ...
  73. def ifftn(
  74. a: ArrayLike,
  75. s: Sequence[int] | None = ...,
  76. axes: Sequence[int] | None = ...,
  77. norm: _NormKind = ...,
  78. out: NDArray[complex128] | None = ...,
  79. ) -> NDArray[complex128]: ...
  80. def rfftn(
  81. a: ArrayLike,
  82. s: Sequence[int] | None = ...,
  83. axes: Sequence[int] | None = ...,
  84. norm: _NormKind = ...,
  85. out: NDArray[complex128] | None = ...,
  86. ) -> NDArray[complex128]: ...
  87. def irfftn(
  88. a: ArrayLike,
  89. s: Sequence[int] | None = ...,
  90. axes: Sequence[int] | None = ...,
  91. norm: _NormKind = ...,
  92. out: NDArray[float64] | None = ...,
  93. ) -> NDArray[float64]: ...
  94. def fft2(
  95. a: ArrayLike,
  96. s: Sequence[int] | None = ...,
  97. axes: Sequence[int] | None = ...,
  98. norm: _NormKind = ...,
  99. out: NDArray[complex128] | None = ...,
  100. ) -> NDArray[complex128]: ...
  101. def ifft2(
  102. a: ArrayLike,
  103. s: Sequence[int] | None = ...,
  104. axes: Sequence[int] | None = ...,
  105. norm: _NormKind = ...,
  106. out: NDArray[complex128] | None = ...,
  107. ) -> NDArray[complex128]: ...
  108. def rfft2(
  109. a: ArrayLike,
  110. s: Sequence[int] | None = ...,
  111. axes: Sequence[int] | None = ...,
  112. norm: _NormKind = ...,
  113. out: NDArray[complex128] | None = ...,
  114. ) -> NDArray[complex128]: ...
  115. def irfft2(
  116. a: ArrayLike,
  117. s: Sequence[int] | None = ...,
  118. axes: Sequence[int] | None = ...,
  119. norm: _NormKind = ...,
  120. out: NDArray[float64] | None = ...,
  121. ) -> NDArray[float64]: ...