Skip to content

Commit 0e725d0

Browse files
jorenhamjorenham
authored
✨ signal: improve resample (#828)
2 parents 2bf8399 + fb3797a commit 0e725d0

File tree

2 files changed

+156
-41
lines changed

2 files changed

+156
-41
lines changed

‎scipy-stubs/signal/_signaltools.pyi‎

Lines changed: 98 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ from collections.abc import Callable
77
from typing import Any, Literal as L, TypeAlias, TypeVar, TypedDict, overload, type_check_only
88

99
import numpy as np
10+
import numpy_typing_compat as nptc
1011
import optype as op
1112
import optype.numpy as onp
1213
import optype.numpy.compat as npc
@@ -54,12 +55,14 @@ __all__ = [
5455
###
5556

5657
_T = TypeVar("_T")
57-
_NumericT = TypeVar("_NumericT", bound=npc.number | np.bool_)
5858
_InexactT = TypeVar("_InexactT", bound=npc.inexact)
59-
_EnvelopeSCT = TypeVar("_EnvelopeSCT", bound=np.float32 | np.float64 | npc.floating80 | npc.complexfloating)
59+
_NumericT = TypeVar("_NumericT", bound=npc.number | np.bool_)
60+
_InexactStandardT = TypeVar("_InexactStandardT", bound=np.float32 | np.float64 | npc.floating80 | npc.complexfloating)
6061
_CoFloat64T = TypeVar("_CoFloat64T", bound=np.float64 | np.float32 | npc.integer)
6162
_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...])
63+
6264
_AnyShapeT = TypeVar("_AnyShapeT", tuple[int], tuple[int, int], tuple[int, int, int], tuple[Any, ...])
65+
_AnyInexact64T = TypeVar("_AnyInexact64T", np.float64, np.complex128)
6366

6467
_Tuple2: TypeAlias = tuple[_T, _T]
6568

@@ -78,8 +81,7 @@ _FilterType: TypeAlias = L["iir", "fir"] | dlti
7881
_F16_64: TypeAlias = np.float64 | np.float32 | np.float16
7982
_C64_128: TypeAlias = np.complex128 | np.complex64
8083

81-
_WindowFuncFloat: TypeAlias = Callable[[onp.Array1D[np.float64]], onp.ToFloat1D]
82-
_WindowFuncComplex: TypeAlias = Callable[[onp.Array1D[np.complex128]], onp.ToFloat1D]
84+
_ToResampleWindow: TypeAlias = Callable[[onp.Array1D[_InexactT]], onp.ToFloat1D] | onp.ToFloat1D | _ToWindow
8385

8486
# workaround for a strange bug in pyright's overlapping overload detection with `numpy<2.1`
8587
_WorkaroundForPyright: TypeAlias = tuple[int] | tuple[Any, ...]
@@ -332,11 +334,11 @@ def fftconvolve(
332334
) -> onp.ArrayND[np.float32, _AnyShapeT]: ...
333335
@overload # generic dtype, generic, shape
334336
def fftconvolve(
335-
in1: onp.ArrayND[_EnvelopeSCT, _AnyShapeT],
336-
in2: onp.ArrayND[_EnvelopeSCT, _AnyShapeT],
337+
in1: onp.ArrayND[_InexactStandardT, _AnyShapeT],
338+
in2: onp.ArrayND[_InexactStandardT, _AnyShapeT],
337339
mode: onp.ConvolveMode = "full",
338340
axes: None = None,
339-
) -> onp.ArrayND[_EnvelopeSCT, _AnyShapeT]: ...
341+
) -> onp.ArrayND[_InexactStandardT, _AnyShapeT]: ...
340342
@overload # ~float64, +float64
341343
def fftconvolve(
342344
in1: onp.ToJustFloat64_ND, in2: onp.ToFloat64_ND, mode: onp.ConvolveMode = "full", axes: AnyShape | None = None
@@ -368,11 +370,11 @@ def oaconvolve(
368370
) -> onp.ArrayND[np.float32, _AnyShapeT]: ...
369371
@overload # generic dtype, generic, shape
370372
def oaconvolve(
371-
in1: onp.ArrayND[_EnvelopeSCT, _AnyShapeT],
372-
in2: onp.ArrayND[_EnvelopeSCT, _AnyShapeT],
373+
in1: onp.ArrayND[_InexactStandardT, _AnyShapeT],
374+
in2: onp.ArrayND[_InexactStandardT, _AnyShapeT],
373375
mode: onp.ConvolveMode = "full",
374376
axes: None = None,
375-
) -> onp.ArrayND[_EnvelopeSCT, _AnyShapeT]: ...
377+
) -> onp.ArrayND[_InexactStandardT, _AnyShapeT]: ...
376378
@overload # ~float64, +float64
377379
def oaconvolve(
378380
in1: onp.ToJustFloat64_ND, in2: onp.ToFloat64_ND, mode: onp.ConvolveMode = "full", axes: AnyShape | None = None
@@ -1019,61 +1021,116 @@ def invresz(
10191021
r: onp.ToComplex1D, p: onp.ToComplex1D, k: onp.ToFloat1D, tol: float = 0.001, rtype: _ResidueType = "avg"
10201022
) -> tuple[onp.Array1D[np.complex128], onp.Array1D[np.complex128]]: ...
10211023

1022-
# TODO(jorenham): improve
1023-
@overload
1024+
# NOTE: We use `_AnyInexact64T` as "free" type parameter, which behaves exactly as
1025+
# the (hypothetical) `AnyOf[np.float64, np.complex128]` gradual type.
1026+
@overload # known dtype, known shape, t=None (default)
10241027
def resample(
1025-
x: onp.ArrayND[_EnvelopeSCT, _AnyShapeT],
1028+
x: nptc.CanArray[_ShapeT, np.dtype[_InexactStandardT]],
10261029
num: int,
10271030
t: None = None,
10281031
axis: int = 0,
1029-
window: _WindowFuncFloat | _WindowFuncComplex | onp.ToFloat1D | _ToWindow | None = None,
1032+
window: _ToResampleWindow[_AnyInexact64T] | None = None,
10301033
domain: _Domain = "time",
1031-
) -> onp.ArrayND[_EnvelopeSCT, _AnyShapeT]: ...
1032-
@overload
1034+
) -> onp.ArrayND[_InexactStandardT, _ShapeT]: ...
1035+
@overload # known dtype, known shape, t=<given>
10331036
def resample(
1034-
x: onp.ToFloatND,
1037+
x: nptc.CanArray[_ShapeT, np.dtype[_InexactStandardT]],
1038+
num: int,
1039+
t: onp.ToFloat1D,
1040+
axis: int = 0,
1041+
window: _ToResampleWindow[_AnyInexact64T] | None = None,
1042+
domain: _Domain = "time",
1043+
) -> tuple[onp.ArrayND[_InexactStandardT, _ShapeT], onp.Array1D[np.float64]]: ...
1044+
@overload # +integer, known shape, t=None (default)
1045+
def resample(
1046+
x: nptc.CanArray[_ShapeT, np.dtype[npc.integer | np.bool_]],
10351047
num: int,
10361048
t: None = None,
10371049
axis: int = 0,
1038-
window: _WindowFuncFloat | onp.ToFloat1D | _ToWindow | None = None,
1050+
window: _ToResampleWindow[np.float64] | None = None,
10391051
domain: _Domain = "time",
1040-
) -> onp.ArrayND[npc.floating]: ...
1041-
@overload
1052+
) -> onp.ArrayND[np.float64, _ShapeT]: ...
1053+
@overload # +integer, known shape, t=<given>
1054+
def resample(
1055+
x: nptc.CanArray[_ShapeT, np.dtype[npc.integer | np.bool_]],
1056+
num: int,
1057+
t: onp.ToFloat1D,
1058+
axis: int = 0,
1059+
window: _ToResampleWindow[np.float64] | None = None,
1060+
domain: _Domain = "time",
1061+
) -> tuple[onp.ArrayND[np.float64, _ShapeT], onp.Array1D[np.float64]]: ...
1062+
@overload # ~float16, known shape, t=None (default)
10421063
def resample(
1043-
x: onp.ToJustComplexND,
1064+
x: nptc.CanArray[_ShapeT, np.dtype[np.float16]],
10441065
num: int,
10451066
t: None = None,
10461067
axis: int = 0,
1047-
window: _WindowFuncComplex | onp.ToFloat1D | _ToWindow | None = None,
1068+
window: _ToResampleWindow[np.float64] | None = None,
10481069
domain: _Domain = "time",
1049-
) -> onp.ArrayND[npc.complexfloating]: ...
1050-
@overload
1070+
) -> onp.ArrayND[np.float32, _ShapeT]: ...
1071+
@overload # ~float16, unknown shape, t=<given>
10511072
def resample(
1052-
x: onp.ArrayND[_EnvelopeSCT, _AnyShapeT],
1073+
x: nptc.CanArray[_ShapeT, np.dtype[np.float16]],
10531074
num: int,
10541075
t: onp.ToFloat1D,
10551076
axis: int = 0,
1056-
window: _WindowFuncFloat | _WindowFuncComplex | onp.ToFloat1D | _ToWindow | None = None,
1077+
window: _ToResampleWindow[np.float64] | None = None,
10571078
domain: _Domain = "time",
1058-
) -> tuple[onp.ArrayND[_EnvelopeSCT, _AnyShapeT], onp.Array1D[npc.floating]]: ...
1059-
@overload
1079+
) -> tuple[onp.ArrayND[np.float32, _ShapeT], onp.Array1D[np.float64]]: ...
1080+
@overload # ~float64 | +integer, unknown shape, t=None (default)
10601081
def resample(
1061-
x: onp.ToFloatND,
1082+
x: onp.SequenceND[float],
1083+
num: int,
1084+
t: None = None,
1085+
axis: int = 0,
1086+
window: _ToResampleWindow[np.float64] | None = None,
1087+
domain: _Domain = "time",
1088+
) -> onp.ArrayND[np.float64]: ...
1089+
@overload # ~float64 | +integer, unknown shape, t=<given>
1090+
def resample(
1091+
x: onp.SequenceND[float],
10621092
num: int,
10631093
t: onp.ToFloat1D,
10641094
axis: int = 0,
1065-
window: _WindowFuncFloat | onp.ToFloat1D | _ToWindow | None = None,
1095+
window: _ToResampleWindow[np.float64] | None = None,
10661096
domain: _Domain = "time",
1067-
) -> tuple[onp.ArrayND[npc.floating], onp.Array1D[npc.floating]]: ...
1068-
@overload
1097+
) -> tuple[onp.ArrayND[np.float64], onp.Array1D[np.float64]]: ...
1098+
@overload # ~complex128, unknown shape, t=None (default)
1099+
def resample(
1100+
x: onp.SequenceND[op.JustComplex | np.complex128],
1101+
num: int,
1102+
t: None = None,
1103+
axis: int = 0,
1104+
window: _ToResampleWindow[np.complex128] | None = None,
1105+
domain: _Domain = "time",
1106+
) -> onp.ArrayND[np.complex128]: ...
1107+
@overload # ~complex128, unknown shape, t=<given>
1108+
def resample(
1109+
x: onp.SequenceND[op.JustComplex | np.complex128],
1110+
num: int,
1111+
t: onp.ToFloat1D,
1112+
axis: int = 0,
1113+
window: _ToResampleWindow[np.complex128] | None = None,
1114+
domain: _Domain = "time",
1115+
) -> tuple[onp.ArrayND[np.complex128], onp.Array1D[np.float64]]: ...
1116+
@overload # unknown dtype, unknown shape, t=None (default)
1117+
def resample(
1118+
x: onp.ToComplexND,
1119+
num: int,
1120+
t: None = None,
1121+
axis: int = 0,
1122+
window: _ToResampleWindow[_AnyInexact64T] | None = None,
1123+
domain: _Domain = "time",
1124+
) -> onp.ArrayND[Any, _WorkaroundForPyright]: ...
1125+
@overload # unknown dtype, unknown shape, t=<given>
10691126
def resample(
1070-
x: onp.ToJustComplexND,
1127+
x: onp.ToComplexND,
10711128
num: int,
10721129
t: onp.ToFloat1D,
10731130
axis: int = 0,
1074-
window: _WindowFuncComplex | onp.ToFloat1D | _ToWindow | None = None,
1131+
window: _ToResampleWindow[_AnyInexact64T] | None = None,
10751132
domain: _Domain = "time",
1076-
) -> tuple[onp.ArrayND[npc.complexfloating], onp.Array1D[npc.floating]]: ...
1133+
) -> tuple[onp.ArrayND[Any, _WorkaroundForPyright], onp.Array1D[np.float64]]: ...
10771134

10781135
# TODO(jorenham): improve
10791136
@overload
@@ -1150,31 +1207,31 @@ def envelope(
11501207
) -> onp.ArrayND[np.float32]: ...
11511208
@overload
11521209
def envelope(
1153-
z: onp.Array1D[_EnvelopeSCT],
1210+
z: onp.Array1D[_InexactStandardT],
11541211
bp_in: tuple[int | None, int | None] = (1, None),
11551212
*,
11561213
n_out: int | None = None,
11571214
squared: bool = False,
11581215
residual: _ResidualKind | None = "lowpass",
11591216
axis: int = -1,
1160-
) -> onp.Array2D[_EnvelopeSCT]: ...
1217+
) -> onp.Array2D[_InexactStandardT]: ...
11611218
@overload
11621219
def envelope(
1163-
z: onp.Array2D[_EnvelopeSCT],
1220+
z: onp.Array2D[_InexactStandardT],
11641221
bp_in: tuple[int | None, int | None] = (1, None),
11651222
*,
11661223
n_out: int | None = None,
11671224
squared: bool = False,
11681225
residual: _ResidualKind | None = "lowpass",
11691226
axis: int = -1,
1170-
) -> onp.Array3D[_EnvelopeSCT]: ...
1227+
) -> onp.Array3D[_InexactStandardT]: ...
11711228
@overload
11721229
def envelope(
1173-
z: onp.ArrayND[_EnvelopeSCT],
1230+
z: onp.ArrayND[_InexactStandardT],
11741231
bp_in: tuple[int | None, int | None] = (1, None),
11751232
*,
11761233
n_out: int | None = None,
11771234
squared: bool = False,
11781235
residual: _ResidualKind | None = "lowpass",
11791236
axis: int = -1,
1180-
) -> onp.ArrayND[_EnvelopeSCT]: ...
1237+
) -> onp.ArrayND[_InexactStandardT]: ...

‎tests/signal/test_resample.pyi‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
# type-tests for `resample` from `signal/_signaltools.pyi`
2+
3+
from typing import assert_type
4+
5+
import numpy as np
6+
import optype.numpy as onp
7+
8+
from scipy.signal import resample
9+
10+
num: int
11+
12+
py_i_1d: list[int]
13+
py_i_2d: list[list[int]]
14+
py_f_1d: list[float]
15+
py_f_2d: list[list[float]]
16+
17+
i8_1d: onp.Array1D[np.int8]
18+
i8_2d: onp.Array2D[np.int8]
19+
20+
f16_1d: onp.Array1D[np.float16]
21+
f16_2d: onp.Array2D[np.float16]
22+
f32_1d: onp.Array1D[np.float32]
23+
f32_2d: onp.Array2D[np.float32]
24+
f64_1d: onp.Array1D[np.float64]
25+
f64_2d: onp.Array2D[np.float64]
26+
f80_1d: onp.Array1D[np.float128]
27+
f80_2d: onp.Array2D[np.float128]
28+
29+
c64_1d: onp.Array1D[np.complex64]
30+
c64_2d: onp.Array2D[np.complex64]
31+
c128_1d: onp.Array1D[np.complex128]
32+
c128_2d: onp.Array2D[np.complex128]
33+
c160_1d: onp.Array1D[np.complex256]
34+
c160_2d: onp.Array2D[np.complex256]
35+
36+
###
37+
38+
assert_type(resample(py_i_1d, num), onp.ArrayND[np.float64])
39+
assert_type(resample(py_f_1d, num), onp.ArrayND[np.float64])
40+
assert_type(resample(i8_1d, num), onp.Array1D[np.float64])
41+
assert_type(resample(f16_1d, num), onp.Array1D[np.float32])
42+
assert_type(resample(f32_1d, num), onp.Array1D[np.float32])
43+
assert_type(resample(f64_1d, num), onp.Array1D[np.float64])
44+
assert_type(resample(f80_1d, num), onp.Array1D[np.float128])
45+
assert_type(resample(c64_1d, num), onp.Array1D[np.complex64])
46+
assert_type(resample(c128_1d, num), onp.Array1D[np.complex128])
47+
assert_type(resample(c160_1d, num), onp.Array1D[np.complex256])
48+
49+
assert_type(resample(py_i_2d, num), onp.ArrayND[np.float64])
50+
assert_type(resample(py_f_2d, num), onp.ArrayND[np.float64])
51+
assert_type(resample(i8_2d, num), onp.Array2D[np.float64])
52+
assert_type(resample(f16_2d, num), onp.Array2D[np.float32])
53+
assert_type(resample(f32_2d, num), onp.Array2D[np.float32])
54+
assert_type(resample(f64_2d, num), onp.Array2D[np.float64])
55+
assert_type(resample(f80_2d, num), onp.Array2D[np.float128])
56+
assert_type(resample(c64_2d, num), onp.Array2D[np.complex64])
57+
assert_type(resample(c128_2d, num), onp.Array2D[np.complex128])
58+
assert_type(resample(c160_2d, num), onp.Array2D[np.complex256])

0 commit comments

Comments
 (0)