8000 Merge pull request #20279 from BvB93/fft · thomasjpfan/numpy@752b276 · GitHub
[go: up one dir, main page]
Skip to content

Commit 752b276

Browse files
charrischarris
authored
Merge pull request numpy#20279 from BvB93/fft
ENH: Add annotations for `np.fft`
2 parents ae4af75 + fb3a205 commit 752b276

File tree

4 files changed

+220
-19
lines changed

4 files changed

+220
-19
lines changed

‎numpy/fft/__init__.pyi

Lines changed: 24 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2,25 +2,30 @@ from typing import Any, List
22

33
from numpy._pytesttester import PytestTester
44

5+
from numpy.fft._pocketfft import (
6+
fft as fft,
7+
ifft as ifft,
8+
rfft as rfft,
9+
irfft as irfft,
10+
hfft as hfft,
11+
ihfft as ihfft,
12+
rfftn as rfftn,
13+
irfftn as irfftn,
14+
rfft2 as rfft2,
15+
irfft2 as irfft2,
16+
fft2 as fft2,
17+
ifft2 as ifft2,
18+
fftn as fftn,
19+
ifftn as ifftn,
20+
)
21+
22+
from numpy.fft.helper import (
23+
fftshift as fftshift,
24+
ifftshift as ifftshift,
25+
fftfreq as fftfreq,
26+
rfftfreq as rfftfreq,
27+
)
28+
529
__all__: List[str]
630
__path__: List[str]
731
test: PytestTester
8-
9-
def fft(a, n=..., axis=..., norm=...): ...
10-
def ifft(a, n=..., axis=..., norm=...): ...
11-
def rfft(a, n=..., axis=..., norm=...): ...
12-
def irfft(a, n=..., axis=..., norm=...): ...
13-
def hfft(a, n=..., axis=..., norm=...): ...
14-
def ihfft(a, n=..., axis=..., norm=...): ...
15-
def fftn(a, s=..., axes=..., norm=...): ...
16-
def ifftn(a, s=..., axes=..., norm=...): ...
17-
def rfftn(a, s=..., axes=..., norm=...): ...
18-
def irfftn(a, s=..., axes=..., norm=...): ...
19-
def fft2(a, s=..., axes=..., norm=...): ...
20-
def ifft2(a, s=..., axes=..., norm=...): ...
21-
def rfft2(a, s=..., axes=..., norm=...): ...
22-
def irfft2(a, s=..., axes=..., norm=...): ...
23-
def fftshift(x, axes=...): ...
24-
def ifftshift(x, axes=...): ...
25-
def fftfreq(n, d=...): ...
26-
def rfftfreq(n, d=...): ...

‎numpy/fft/_pocketfft.pyi

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

‎numpy/fft/helper.pyi

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
from typing import List, Any, TypeVar, overload
2+
3+
from numpy import generic, dtype, integer, floating, complexfloating
4+
from numpy.typing import (
5+
NDArray,
6+
ArrayLike,
7+
_ShapeLike,
8+
_SupportsArray,
9+
_FiniteNestedSequence,
10+
_ArrayLikeFloat_co,
11+
_ArrayLikeComplex_co,
12+
)
13+
14+
_SCT = TypeVar("_SCT", bound=generic)
15+
16+
_ArrayLike = _FiniteNestedSequence[_SupportsArray[dtype[_SCT]]]
17+
18+
__all__: List[str]
19+
20+
@overload
21+
def fftshift(x: _ArrayLike[_SCT], axes: None | _ShapeLike = ...) -> NDArray[_SCT]: ...
22+
@overload
23+
def fftshift(x: ArrayLike, axes: None | _ShapeLike = ...) -> NDArray[Any]: ...
24+
25+
@overload
26+
def ifftshift(x: _ArrayLike[_SCT], axes: None | _ShapeLike = ...) -> NDArray[_SCT]: ...
27+
@overload
28+
def ifftshift(x: ArrayLike, axes: None | _ShapeLike = ...) -> NDArray[Any]: ...
29+
30+
@overload
31+
def fftfreq(
32+
n: int | integer[Any],
33+
d: _ArrayLikeFloat_co,
34+
) -> NDArray[floating[Any]]: ...
35+
@overload
36+
def fftfreq(
37+
n: int | integer[Any],
38+
d: _ArrayLikeComplex_co,
39+
) -> NDArray[complexfloating[Any, Any]]: ...
40+
41+
@overload
42+
def rfftfreq(
43+
n: int | integer[Any],
44+
d: _ArrayLikeFloat_co,
45+
) -> NDArray[floating[Any]]: ...
46+
@overload
47+
def rfftfreq(
48+
n: int | integer[Any],
49+
d: _ArrayLikeComplex_co,
50+
) -> NDArray[complexfloating[Any, Any]]: ...

‎numpy/typing/tests/data/reveal/fft.pyi

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import numpy as np
2+
import numpy.typing as npt
3+
4+
AR_f8: npt.NDArray[np.float64]
5+
AR_c16: npt.NDArray[np.complex128]
6+
AR_LIKE_f8: list[float]
7+
8+
reveal_type(np.fft.fftshift(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
9+
reveal_type(np.fft.fftshift(AR_LIKE_f8, axes=0)) # E: numpy.ndarray[Any, numpy.dtype[Any]]
10+
11+
reveal_type(np.fft.ifftshift(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
12+
reveal_type(np.fft.ifftshift(AR_LIKE_f8, axes=0)) # E: numpy.ndarray[Any, numpy.dtype[Any]]
13+
14+
reveal_type(np.fft.fftfreq(5, AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
15+
reveal_type(np.fft.fftfreq(np.int64(), AR_c16)) # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
16+
17+
reveal_type(np.fft.fftfreq(5, AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
18+
reveal_type(np.fft.fftfreq(np.int64(), AR_c16)) # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
19+
20+
reveal_type(np.fft.fft(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
21+
reveal_type(np.fft.ifft(AR_f8, axis=1)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
22+
reveal_type(np.fft.rfft(AR_f8, n=None)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
23+
reveal_type(np.fft.irfft(AR_f8, norm="ortho")) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
24+
reveal_type(np.fft.hfft(AR_f8, n=2)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
25+
reveal_type(np.fft.ihfft(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
26+
27+
reveal_type(np.fft.fftn(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
28+
reveal_type(np.fft.ifftn(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
29+
reveal_type(np.fft.rfftn(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
30+
reveal_type(np.fft.irfftn(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
31+
32+
reveal_type(np.fft.rfft2(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
33+
reveal_type(np.fft.ifft2(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
34+
reveal_type(np.fft.fft2(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
35+
reveal_type(np.fft.irfft2(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]

0 commit comments

Comments
 (0)
0