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