Skip to content

Commit a03b355

Browse files
jorenhamjorenham
authored
Merge pull request #157 from pavyamsiri/improve/signal/spectral
`signal`: Add type stubs to `_spectral_py.pyi`.
2 parents 818150c + fb5c85d commit a03b355

File tree

2 files changed

+261
-86
lines changed

2 files changed

+261
-86
lines changed

‎scipy-stubs/signal/_spectral_py.pyi‎

Lines changed: 196 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -1,97 +1,207 @@
1-
from typing import Literal
1+
from collections.abc import Callable
2+
from typing import Literal, TypeAlias, overload
3+
from typing_extensions import Unpack
24

3-
from scipy._typing import Untyped, UntypedCallable
5+
import numpy as np
6+
import numpy.typing as npt
7+
import optype as op
8+
from numpy._typing import _ArrayLikeFloat_co, _ArrayLikeNumber_co
9+
from scipy._typing import AnyInt, AnyReal
10+
from .windows._windows import _Window, _WindowNeedsParams
411

512
__all__ = ["check_COLA", "check_NOLA", "coherence", "csd", "istft", "lombscargle", "periodogram", "spectrogram", "stft", "welch"]
613

7-
def lombscargle(x: Untyped, y: Untyped, freqs: Untyped, precenter: bool = False, normalize: bool = False) -> Untyped: ...
14+
_Array_f8: TypeAlias = npt.NDArray[np.float64]
15+
_Array_f8_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]]
16+
_ArrayFloat: TypeAlias = npt.NDArray[np.float32 | np.float64 | np.longdouble]
17+
_ArrayComplex: TypeAlias = npt.NDArray[np.complex64 | np.complex128 | np.clongdouble]
18+
19+
_GetWindowArgument: TypeAlias = _Window | tuple[_Window | _WindowNeedsParams, Unpack[tuple[object, ...]]]
20+
_WindowLike: TypeAlias = _GetWindowArgument | _ArrayLikeFloat_co
21+
_Detrend: TypeAlias = Literal["literal", "constant", False] | Callable[[npt.NDArray[np.generic]], npt.NDArray[np.generic]]
22+
_Scaling: TypeAlias = Literal["density", "spectrum"]
23+
_LegacyScaling: TypeAlias = Literal["psd", "spectrum"]
24+
_Average: TypeAlias = Literal["mean", "median"]
25+
_Boundary: TypeAlias = Literal["even", "odd", "constant", "zeros"] | None
26+
27+
def lombscargle(
28+
x: _ArrayLikeFloat_co,
29+
y: _ArrayLikeFloat_co,
30+
freqs: _ArrayLikeFloat_co,
31+
precenter: op.CanBool = False,
32+
normalize: op.CanBool = False,
33+
) -> _Array_f8_1d: ...
834
def periodogram(
9-
x: Untyped,
10-
fs: float = 1.0,
11-
window: str = "boxcar",
12-
nfft: int | None = None,
13-
detrend: str | Literal[False] | UntypedCallable = "constant",
14-
return_onesided: bool = True,
15-
scaling: str = "density",
16-
axis: int = -1,
17-
) -> Untyped: ...
35+
x: _ArrayLikeNumber_co,
36+
fs: AnyReal = 1.0,
37+
window: _WindowLike | None = "boxcar",
38+
nfft: AnyInt | None = None,
39+
detrend: _Detrend = "constant",
40+
return_onesided: op.CanBool = True,
41+
scaling: _Scaling = "density",
42+
axis: op.CanIndex = -1,
43+
) -> tuple[_Array_f8, _ArrayFloat]: ...
1844
def welch(
19-
x: Untyped,
20-
fs: float = 1.0,
21-
window: str = "hann",
22-
nperseg: int | None = None,
23-
noverlap: int | None = None,
24-
nfft: int | None = None,
25-
detrend: str | Literal[False] | UntypedCallable = "constant",
26-
return_onesided: bool = True,
27-
scaling: str = "density",
28-
axis: int = -1,
29-
average: str = "mean",
30-
) -> Untyped: ...
45+
x: _ArrayLikeNumber_co,
46+
fs: AnyReal = 1.0,
47+
window: _WindowLike = "hann",
48+
nperseg: AnyInt | None = None,
49+
noverlap: AnyInt | None = None,
50+
nfft: AnyInt | None = None,
51+
detrend: _Detrend = "constant",
52+
return_onesided: op.CanBool = True,
53+
scaling: _Scaling = "density",
54+
axis: op.CanIndex = -1,
55+
average: _Average = "mean",
56+
) -> tuple[_Array_f8, _ArrayFloat]: ...
3157
def csd(
32-
x: Untyped,
33-
y: Untyped,
34-
fs: float = 1.0,
35-
window: str = "hann",
36-
nperseg: int | None = None,
37-
noverlap: int | None = None,
38-
nfft: int | None = None,
39-
detrend: str | Literal[False] | UntypedCallable = "constant",
40-
return_onesided: bool = True,
41-
scaling: str = "density",
42-
axis: int = -1,
43-
average: str = "mean",
44-
) -> Untyped: ...
58+
x: _ArrayLikeNumber_co,
59+
y: _ArrayLikeNumber_co,
60+
fs: AnyReal = 1.0,
61+
window: _WindowLike = "hann",
62+
nperseg: AnyInt | None = None,
63+
noverlap: AnyInt | None = None,
64+
nfft: AnyInt | None = None,
65+
detrend: _Detrend = "constant",
66+
return_onesided: op.CanBool = True,
67+
scaling: _Scaling = "density",
68+
axis: op.CanIndex = -1,
69+
average: _Average = "mean",
70+
) -> tuple[_Array_f8, _ArrayComplex]: ...
71+
72+
#
73+
@overload
74+
# non-complex mode (positional and keyword)
75+
def spectrogram(
76+
x: _ArrayLikeNumber_co,
77+
fs: AnyReal = 1.0,
78+
window: _WindowLike = ("tukey", 0.25),
79+
nperseg: AnyInt | None = None,
80+
noverlap: AnyInt | None = None,
81+
nfft: AnyInt | None = None,
82+
detrend: _Detrend = "constant",
83+
return_onesided: op.CanBool = True,
84+
scaling: _Scaling = "density",
85+
axis: op.CanIndex = -1,
86+
mode: Literal["psd", "magnitude", "angle", "phase"] = "psd",
87+
) -> tuple[_Array_f8, _Array_f8, _ArrayFloat]: ...
88+
@overload
89+
# complex mode (positional)
90+
def spectrogram(
91+
x: _ArrayLikeNumber_co,
92+
fs: AnyReal,
93+
window: _WindowLike,
94+
nperseg: AnyInt | None,
95+
noverlap: AnyInt | None,
96+
nfft: AnyInt | None,
97+
detrend: _Detrend,
98+
return_onesided: op.CanBool,
99+
scaling: _Scaling,
100+
axis: op.CanIndex,
101+
mode: Literal["complex"],
102+
) -> tuple[_Array_f8, _Array_f8, _ArrayComplex]: ...
103+
@overload
104+
# complex mode (keyword)
45105
def spectrogram(
46-
x: Untyped,
47-
fs: float = 1.0,
48-
window: Untyped = ("tukey", 0.25),
49-
nperseg: int | None = None,
50-
noverlap: int | None = None,
51-
nfft: int | None = None,
52-
detrend: str | Literal[False] | UntypedCallable = "constant",
53-
return_onesided: bool = True,
54-
scaling: str = "density",
55-
axis: int = -1,
56-
mode: str = "psd",
57-
) -> Untyped: ...
58-
def check_COLA(window: Untyped, nperseg: int, noverlap: int, tol: float = 1e-10) -> Untyped: ...
59-
def check_NOLA(window: Untyped, nperseg: int, noverlap: int, tol: float = 1e-10) -> Untyped: ...
106+
x: _ArrayLikeNumber_co,
107+
fs: AnyReal = 1.0,
108+
window: _WindowLike = ("tukey", 0.25),
109+
nperseg: AnyInt | None = None,
110+
noverlap: AnyInt | None = None,
111+
nfft: AnyInt | None = None,
112+
detrend: _Detrend = "constant",
113+
return_onesided: op.CanBool = True,
114+
scaling: _Scaling = "density",
115+
axis: op.CanIndex = -1,
116+
*,
117+
mode: Literal["complex"],
118+
) -> tuple[_Array_f8, _Array_f8, _ArrayComplex]: ...
119+
120+
#
121+
def check_COLA(
122+
window: _WindowLike,
123+
nperseg: AnyInt,
124+
noverlap: AnyInt,
125+
tol: AnyReal = 1e-10,
126+
) -> np.bool_: ...
127+
def check_NOLA(
128+
window: _WindowLike,
129+
nperseg: AnyInt,
130+
noverlap: AnyInt,
131+
tol: AnyReal = 1e-10,
132+
) -> np.bool_: ...
60133
def stft(
61-
x: Untyped,
62-
fs: float = 1.0,
63-
window: str = "hann",
64-
nperseg: int = 256,
65-
noverlap: int | None = None,
66-
nfft: int | None = None,
67-
detrend: bool = False,
68-
return_onesided: bool = True,
69-
boundary: str = "zeros",
70-
padded: bool = True,
71-
axis: int = -1,
72-
scaling: str = "spectrum",
73-
) -> Untyped: ...
134+
x: _ArrayLikeNumber_co,
135+
fs: AnyReal = 1.0,
136+
window: _WindowLike = "hann",
137+
nperseg: AnyInt = 256,
138+
noverlap: AnyInt | None = None,
139+
nfft: AnyInt | None = None,
140+
detrend: _Detrend = False,
141+
return_onesided: op.CanBool = True,
142+
boundary: _Boundary = "zeros",
143+
padded: op.CanBool = True,
144+
axis: op.CanIndex = -1,
145+
scaling: _LegacyScaling = "spectrum",
146+
) -> tuple[_Array_f8, _Array_f8, _ArrayComplex]: ...
147+
148+
#
149+
@overload
150+
# input_onesided is `True`
151+
def istft(
152+
Zxx: _ArrayLikeNumber_co,
153+
fs: AnyReal = 1.0,
154+
window: _WindowLike = "hann",
155+
nperseg: AnyInt | None = None,
156+
noverlap: AnyInt | None = None,
157+
nfft: AnyInt | None = None,
158+
input_onesided: Literal[True, 1] = True,
159+
boundary: op.CanBool = True,
160+
time_axis: op.CanIndex = -1,
161+
freq_axis: op.CanIndex = -2,
162+
scaling: _LegacyScaling = "spectrum",
163+
) -> tuple[_Array_f8, _ArrayFloat]: ...
164+
@overload
165+
# input_onesided is `False` (positional)
74166
def istft(
75-
Zxx: Untyped,
76-
fs: float = 1.0,
77-
window: str = "hann",
78-
nperseg: int | None = None,
79-
noverlap: int | None = None,
80-
nfft: int | None = None,
81-
input_onesided: bool = True,
82-
boundary: bool = True,
83-
time_axis: int = -1,
84-
freq_axis: int = -2,
85-
scaling: str = "spectrum",
86-
) -> Untyped: ...
167+
Zxx: _ArrayLikeNumber_co,
168+
fs: AnyReal,
169+
window: _WindowLike,
170+
nperseg: AnyInt | None,
171+
noverlap: AnyInt | None,
172+
nfft: AnyInt | None,
173+
input_onesided: Literal[False, 0],
174+
boundary: op.CanBool = True,
175+
time_axis: op.CanIndex = -1,
176+
freq_axis: op.CanIndex = -2,
177+
scaling: _LegacyScaling = "spectrum",
178+
) -> tuple[_Array_f8, _ArrayComplex]: ...
179+
@overload
180+
# input_onesided is `False` (keyword)
181+
def istft(
182+
Zxx: _ArrayLikeNumber_co,
183+
fs: AnyReal = 1.0,
184+
window: _WindowLike = "hann",
185+
nperseg: AnyInt | None = None,
186+
noverlap: AnyInt | None = None,
187+
nfft: AnyInt | None = None,
188+
*,
189+
input_onesided: Literal[False, 0],
190+
boundary: op.CanBool = True,
191+
time_axis: op.CanIndex = -1,
192+
freq_axis: op.CanIndex = -2,
193+
scaling: _LegacyScaling = "spectrum",
194+
) -> tuple[_Array_f8, _ArrayComplex]: ...
195+
196+
#
87197
def coherence(
88-
x: Untyped,
89-
y: Untyped,
90-
fs: float = 1.0,
91-
window: str = "hann",
92-
nperseg: int | None = None,
93-
noverlap: int | None = None,
94-
nfft: int | None = None,
95-
detrend: str | Literal[False] | UntypedCallable = "constant",
96-
axis: int = -1,
97-
) -> Untyped: ...
198+
x: _ArrayLikeNumber_co,
199+
y: _ArrayLikeNumber_co,
200+
fs: AnyReal = 1.0,
201+
window: _WindowLike = "hann",
202+
nperseg: AnyInt | None = None,
203+
noverlap: AnyInt | None = None,
204+
nfft: AnyInt | None = None,
205+
detrend: _Detrend = "constant",
206+
axis: op.CanIndex = -1,
207+
) -> tuple[_Array_f8, _ArrayFloat]: ...

‎tests/signal/test_spectral.pyi‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
from typing import Literal, TypeAlias
2+
from typing_extensions import assert_type
3+
4+
import numpy as np
5+
import numpy.typing as npt
6+
import optype.numpy as onpt
7+
from scipy.signal import istft, spectrogram
8+
9+
_Array_f8: TypeAlias = npt.NDArray[np.float64]
10+
_ArrayFloat: TypeAlias = npt.NDArray[np.float32 | np.float64 | np.longdouble]
11+
_ArrayComplex: TypeAlias = npt.NDArray[np.complex64 | np.complex128 | np.clongdouble]
12+
13+
array_f8_1d: onpt.Array[tuple[Literal[256]], np.float64]
14+
array_c16_1d: onpt.Array[tuple[Literal[256]], np.complex128]
15+
spectrogram_mode_real: Literal["psd", "magnitude", "angle", "phase"]
16+
17+
# test spectrogram function overloads
18+
assert_type(spectrogram(array_f8_1d), tuple[_Array_f8, _Array_f8, _ArrayFloat])
19+
assert_type(spectrogram(array_f8_1d, mode=spectrogram_mode_real), tuple[_Array_f8, _Array_f8, _ArrayFloat])
20+
assert_type(spectrogram(array_f8_1d, mode="complex"), tuple[_Array_f8, _Array_f8, _ArrayComplex])
21+
assert_type(
22+
spectrogram(array_f8_1d, 1.0, ("tukey", 2.5), None, None, None, "constant", True, "density", -1, "complex"),
23+
tuple[_Array_f8, _Array_f8, _ArrayComplex],
24+
)
25+
26+
# test isft function overloads
27+
assert_type(istft(array_c16_1d), tuple[_Array_f8, _ArrayFloat])
28+
assert_type(istft(array_c16_1d, input_onesided=True), tuple[_Array_f8, _ArrayFloat])
29+
assert_type(istft(array_c16_1d, 1.0, "hann", 256, 128, 256, False), tuple[_Array_f8, _ArrayComplex])
30+
assert_type(
31+
istft(array_c16_1d, input_onesided=False, fs=1.0, window="hann", nperseg=256, noverlap=128, nfft=256),
32+
tuple[_Array_f8, _ArrayComplex],
33+
)
34+
assert_type(
35+
istft(
36+
array_c16_1d,
37+
fs=2.0,
38+
window=("tukey", 0.25),
39+
nperseg=256,
40+
noverlap=128,
41+
nfft=256,
42+
input_onesided=True,
43+
boundary=False,
44+
time_axis=-1,
45+
freq_axis=0,
46+
scaling="spectrum",
47+
),
48+
tuple[_Array_f8, _ArrayFloat],
49+
)
50+
assert_type(
51+
istft(
52+
array_c16_1d,
53+
fs=2.0,
54+
window=("tukey", 0.25),
55+
nperseg=256,
56+
noverlap=128,
57+
nfft=256,
58+
input_onesided=False,
59+
boundary=False,
60+
time_axis=0,
61+
freq_axis=1,
62+
scaling="spectrum",
63+
),
64+
tuple[_Array_f8, _ArrayComplex],
65+
)

0 commit comments

Comments
 (0)