Skip to content

Commit 2dccd6a

Browse files
jorenhamjorenham
committed
✨ fft: improved Nd fft annotations
1 parent f89195a commit 2dccd6a

File tree

1 file changed

+293
-15
lines changed

1 file changed

+293
-15
lines changed

‎scipy-stubs/fft/_basic.pyi

Lines changed: 293 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,6 @@ _ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...])
1313
_Norm: TypeAlias = Literal["backward", "ortho", "forward"]
1414
_Unused: TypeAlias = Never # not used by scipy
1515

16-
_FloatingND: TypeAlias = onp.ArrayND[np.float32 | np.float64 | np.longdouble]
17-
_ComplexND: TypeAlias = onp.ArrayND[npc.complexfloating]
18-
1916
_AsFloat32: TypeAlias = onp.CanArray[_ShapeT, np.dtype[npc.floating32]]
2017
_AsFloat64: TypeAlias = onp.CanArray[_ShapeT, np.dtype[np.bool_ | npc.integer | npc.floating64]]
2118
_AsFloat80: TypeAlias = onp.CanArray[_ShapeT, np.dtype[np.longdouble]]
@@ -853,63 +850,344 @@ def ihfft2(
853850
###
854851
# N-D
855852

853+
# keep in sync with `ifftn`
854+
@overload
855+
def fftn( # type: ignore[overload-overlap]
856+
x: _AsComplex64[_ShapeT],
857+
s: onp.ToJustIntND | None = None,
858+
axes: AnyShape | None = None,
859+
norm: _Norm | None = None,
860+
overwrite_x: onp.ToBool = False,
861+
workers: int | None = None,
862+
*,
863+
plan: _Unused | None = None,
864+
) -> onp.Array[_ShapeT, np.complex64]: ...
865+
@overload
866+
def fftn(
867+
x: _AsComplex128[_ShapeT],
868+
s: onp.ToJustIntND | None = None,
869+
axes: AnyShape | None = None,
870+
norm: _Norm | None = None,
871+
overwrite_x: onp.ToBool = False,
872+
workers: int | None = None,
873+
*,
874+
plan: _Unused | None = None,
875+
) -> onp.Array[_ShapeT, np.complex128]: ...
876+
@overload
877+
def fftn(
878+
x: _AsComplex160[_ShapeT],
879+
s: onp.ToJustIntND | None = None,
880+
axes: AnyShape | None = None,
881+
norm: _Norm | None = None,
882+
overwrite_x: onp.ToBool = False,
883+
workers: int | None = None,
884+
*,
885+
plan: _Unused | None = None,
886+
) -> onp.Array[_ShapeT, np.clongdouble]: ...
887+
@overload
888+
def fftn(
889+
x: _ToComplex128_ND,
890+
s: onp.ToJustIntND | None = None,
891+
axes: AnyShape | None = None,
892+
norm: _Norm | None = None,
893+
overwrite_x: onp.ToBool = False,
894+
workers: int | None = None,
895+
*,
896+
plan: _Unused | None = None,
897+
) -> onp.ArrayND[np.complex128]: ...
898+
@overload
856899
def fftn(
857900
x: onp.ToComplexND,
858-
s: onp.ToIntND | None = None,
901+
s: onp.ToJustIntND | None = None,
902+
axes: AnyShape | None = None,
903+
norm: _Norm | None = None,
904+
overwrite_x: onp.ToBool = False,
905+
workers: int | None = None,
906+
*,
907+
plan: _Unused | None = None,
908+
) -> onp.ArrayND[np.complex128 | Any]: ...
909+
910+
# keep in sync with `fftn`
911+
@overload
912+
def ifftn( # type: ignore[overload-overlap]
913+
x: _AsComplex64[_ShapeT],
914+
s: onp.ToJustIntND | None = None,
915+
axes: AnyShape | None = None,
916+
norm: _Norm | None = None,
917+
overwrite_x: onp.ToBool = False,
918+
workers: int | None = None,
919+
*,
920+
plan: _Unused | None = None,
921+
) -> onp.Array[_ShapeT, np.complex64]: ...
922+
@overload
923+
def ifftn(
924+
x: _AsComplex128[_ShapeT],
925+
s: onp.ToJustIntND | None = None,
926+
axes: AnyShape | None = None,
927+
norm: _Norm | None = None,
928+
overwrite_x: onp.ToBool = False,
929+
workers: int | None = None,
930+
*,
931+
plan: _Unused | None = None,
932+
) -> onp.Array[_ShapeT, np.complex128]: ...
933+
@overload
934+
def ifftn(
935+
x: _AsComplex160[_ShapeT],
936+
s: onp.ToJustIntND | None = None,
937+
axes: AnyShape | None = None,
938+
norm: _Norm | None = None,
939+
overwrite_x: onp.ToBool = False,
940+
workers: int | None = None,
941+
*,
942+
plan: _Unused | None = None,
943+
) -> onp.Array[_ShapeT, np.clongdouble]: ...
944+
@overload
945+
def ifftn(
946+
x: _ToComplex128_ND,
947+
s: onp.ToJustIntND | None = None,
859948
axes: AnyShape | None = None,
860949
norm: _Norm | None = None,
861950
overwrite_x: onp.ToBool = False,
862951
workers: int | None = None,
863952
*,
864953
plan: _Unused | None = None,
865-
) -> _ComplexND: ...
954+
) -> onp.ArrayND[np.complex128]: ...
955+
@overload
866956
def ifftn(
867957
x: onp.ToComplexND,
868-
s: onp.ToIntND | None = None,
958+
s: onp.ToJustIntND | None = None,
959+
axes: AnyShape | None = None,
960+
norm: _Norm | None = None,
961+
overwrite_x: onp.ToBool = False,
962+
workers: int | None = None,
963+
*,
964+
plan: _Unused | None = None,
965+
) -> onp.ArrayND[np.complex128 | Any]: ...
966+
967+
# keep in sync with `ihfftn`
968+
@overload
969+
def rfftn( # type: ignore[overload-overlap]
970+
x: _AsFloat32[_ShapeT],
971+
s: onp.ToJustIntND | None = None,
869972
axes: AnyShape | None = None,
870973
norm: _Norm | None = None,
871974
overwrite_x: onp.ToBool = False,
872975
workers: int | None = None,
873976
*,
874977
plan: _Unused | None = None,
875-
) -> _ComplexND: ...
978+
) -> onp.Array[_ShapeT, np.complex64]: ...
979+
@overload
980+
def rfftn(
981+
x: _AsFloat64[_ShapeT],
982+
s: onp.ToJustIntND | None = None,
983+
axes: AnyShape | None = None,
984+
norm: _Norm | None = None,
985+
overwrite_x: onp.ToBool = False,
986+
workers: int | None = None,
987+
*,
988+
plan: _Unused | None = None,
989+
) -> onp.Array[_ShapeT, np.complex128]: ...
990+
@overload
991+
def rfftn(
992+
x: _AsFloat80[_ShapeT],
993+
s: onp.ToJustIntND | None = None,
994+
axes: AnyShape | None = None,
995+
norm: _Norm | None = None,
996+
overwrite_x: onp.ToBool = False,
997+
workers: int | None = None,
998+
*,
999+
plan: _Unused | None = None,
1000+
) -> onp.Array[_ShapeT, np.clongdouble]: ...
1001+
@overload
1002+
def rfftn(
1003+
x: _ToFloat64_ND,
1004+
s: onp.ToJustIntND | None = None,
1005+
axes: AnyShape | None = None,
1006+
norm: _Norm | None = None,
1007+
overwrite_x: onp.ToBool = False,
1008+
workers: int | None = None,
1009+
*,
1010+
plan: _Unused | None = None,
1011+
) -> onp.ArrayND[np.complex128]: ...
1012+
@overload
8761013
def rfftn(
8771014
x: onp.ToFloatND,
878-
s: onp.ToIntND | None = None,
1015+
s: onp.ToJustIntND | None = None,
1016+
axes: AnyShape | None = None,
1017+
norm: _Norm | None = None,
1018+
overwrite_x: onp.ToBool = False,
1019+
workers: int | None = None,
1020+
*,
1021+
plan: _Unused | None = None,
1022+
) -> onp.ArrayND[np.complex128 | Any]: ...
1023+
1024+
# keep in sync with `hfftn`
1025+
@overload
1026+
def irfftn( # type: ignore[overload-overlap]
1027+
x: _AsComplex64[_ShapeT],
1028+
s: onp.ToJustIntND | None = None,
1029+
axes: AnyShape | None = None,
1030+
norm: _Norm | None = None,
1031+
overwrite_x: onp.ToBool = False,
1032+
workers: int | None = None,
1033+
*,
1034+
plan: _Unused | None = None,
1035+
) -> onp.Array[_ShapeT, np.float32]: ...
1036+
@overload
1037+
def irfftn(
1038+
x: _AsComplex128[_ShapeT],
1039+
s: onp.ToJustIntND | None = None,
1040+
axes: AnyShape | None = None,
1041+
norm: _Norm | None = None,
1042+
overwrite_x: onp.ToBool = False,
1043+
workers: int | None = None,
1044+
*,
1045+
plan: _Unused | None = None,
1046+
) -> onp.Array[_ShapeT, np.float64]: ...
1047+
@overload
1048+
def irfftn(
1049+
x: _AsComplex160[_ShapeT],
1050+
s: onp.ToJustIntND | None = None,
8791051
axes: AnyShape | None = None,
8801052
norm: _Norm | None = None,
8811053
overwrite_x: onp.ToBool = False,
8821054
workers: int | None = None,
8831055
*,
8841056
plan: _Unused | None = None,
885-
) -> _ComplexND: ...
1057+
) -> onp.Array[_ShapeT, np.longdouble]: ...
1058+
@overload
1059+
def irfftn(
1060+
x: _ToComplex128_ND,
1061+
s: onp.ToJustIntND | None = None,
1062+
axes: AnyShape | None = None,
1063+
norm: _Norm | None = None,
1064+
overwrite_x: onp.ToBool = False,
1065+
workers: int | None = None,
1066+
*,
1067+
plan: _Unused | None = None,
1068+
) -> onp.ArrayND[np.float64]: ...
1069+
@overload
8861070
def irfftn(
8871071
x: onp.ToComplexND,
888-
s: onp.ToIntND | None = None,
1072+
s: onp.ToJustIntND | None = None,
1073+
axes: AnyShape | None = None,
1074+
norm: _Norm | None = None,
1075+
overwrite_x: onp.ToBool = False,
1076+
workers: int | None = None,
1077+
*,
1078+
plan: _Unused | None = None,
1079+
) -> onp.ArrayND[np.float64 | Any]: ...
1080+
1081+
# keep in sync with `irfftn`
1082+
@overload
1083+
def hfftn( # type: ignore[overload-overlap]
1084+
x: _AsComplex64[_ShapeT],
1085+
s: onp.ToJustIntND | None = None,
1086+
axes: AnyShape | None = None,
1087+
norm: _Norm | None = None,
1088+
overwrite_x: onp.ToBool = False,
1089+
workers: int | None = None,
1090+
*,
1091+
plan: _Unused | None = None,
1092+
) -> onp.Array[_ShapeT, np.float32]: ...
1093+
@overload
1094+
def hfftn(
1095+
x: _AsComplex128[_ShapeT],
1096+
s: onp.ToJustIntND | None = None,
1097+
axes: AnyShape | None = None,
1098+
norm: _Norm | None = None,
1099+
overwrite_x: onp.ToBool = False,
1100+
workers: int | None = None,
1101+
*,
1102+
plan: _Unused | None = None,
1103+
) -> onp.Array[_ShapeT, np.float64]: ...
1104+
@overload
1105+
def hfftn(
1106+
x: _AsComplex160[_ShapeT],
1107+
s: onp.ToJustIntND | None = None,
1108+
axes: AnyShape | None = None,
1109+
norm: _Norm | None = None,
1110+
overwrite_x: onp.ToBool = False,
1111+
workers: int | None = None,
1112+
*,
1113+
plan: _Unused | None = None,
1114+
) -> onp.Array[_ShapeT, np.longdouble]: ...
1115+
@overload
1116+
def hfftn(
1117+
x: _ToComplex128_ND,
1118+
s: onp.ToJustIntND | None = None,
8891119
axes: AnyShape | None = None,
8901120
norm: _Norm | None = None,
8911121
overwrite_x: onp.ToBool = False,
8921122
workers: int | None = None,
8931123
*,
8941124
plan: _Unused | None = None,
895-
) -> _FloatingND: ...
1125+
) -> onp.ArrayND[np.float64]: ...
1126+
@overload
8961127
def hfftn(
8971128
x: onp.ToComplexND,
898-
s: onp.ToIntND | None = None,
1129+
s: onp.ToJustIntND | None = None,
1130+
axes: AnyShape | None = None,
1131+
norm: _Norm | None = None,
1132+
overwrite_x: onp.ToBool = False,
1133+
workers: int | None = None,
1134+
*,
1135+
plan: _Unused | None = None,
1136+
) -> onp.ArrayND[np.float64 | Any]: ...
1137+
1138+
# keep in sync with `rfftn`
1139+
@overload
1140+
def ihfftn( # type: ignore[overload-overlap]
1141+
x: _AsFloat32[_ShapeT],
1142+
s: onp.ToJustIntND | None = None,
8991143
axes: AnyShape | None = None,
9001144
norm: _Norm | None = None,
9011145
overwrite_x: onp.ToBool = False,
9021146
workers: int | None = None,
9031147
*,
9041148
plan: _Unused | None = None,
905-
) -> _FloatingND: ...
1149+
) -> onp.Array[_ShapeT, np.complex64]: ...
1150+
@overload
1151+
def ihfftn(
1152+
x: _AsFloat64[_ShapeT],
1153+
s: onp.ToJustIntND | None = None,
1154+
axes: AnyShape | None = None,
1155+
norm: _Norm | None = None,
1156+
overwrite_x: onp.ToBool = False,
1157+
workers: int | None = None,
1158+
*,
1159+
plan: _Unused | None = None,
1160+
) -> onp.Array[_ShapeT, np.complex128]: ...
1161+
@overload
1162+
def ihfftn(
1163+
x: _AsFloat80[_ShapeT],
1164+
s: onp.ToJustIntND | None = None,
1165+
axes: AnyShape | None = None,
1166+
norm: _Norm | None = None,
1167+
overwrite_x: onp.ToBool = False,
1168+
workers: int | None = None,
1169+
*,
1170+
plan: _Unused | None = None,
1171+
) -> onp.Array[_ShapeT, np.clongdouble]: ...
1172+
@overload
1173+
def ihfftn(
1174+
x: _ToFloat64_ND,
1175+
s: onp.ToJustIntND | None = None,
1176+
axes: AnyShape | None = None,
1177+
norm: _Norm | None = None,
1178+
overwrite_x: onp.ToBool = False,
1179+
workers: int | None = None,
1180+
*,
1181+
plan: _Unused | None = None,
1182+
) -> onp.ArrayND[np.complex128]: ...
1183+
@overload
9061184
def ihfftn(
9071185
x: onp.ToFloatND,
908-
s: onp.ToIntND | None = None,
1186+
s: onp.ToJustIntND | None = None,
9091187
axes: AnyShape | None = None,
9101188
norm: _Norm | None = None,
9111189
overwrite_x: onp.ToBool = False,
9121190
workers: int | None = None,
9131191
*,
9141192
plan: _Unused | None = None,
915-
) -> _ComplexND: ...
1193+
) -> onp.ArrayND[np.complex128 | Any]: ...

0 commit comments

Comments
 (0)