_pocketfft.pyi 3.1 KB

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