-
-
Notifications
You must be signed in to change notification settings - Fork 7
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #157 from pavyamsiri/improve/signal/spectral
`signal`: Add type stubs to `_spectral_py.pyi`.
- Loading branch information
Showing
2 changed files
with
261 additions
and
86 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,97 +1,207 @@ | ||
| from typing import Literal | ||
| from collections.abc import Callable | ||
| from typing import Literal, TypeAlias, overload | ||
| from typing_extensions import Unpack | ||
|
|
||
| from scipy._typing import Untyped, UntypedCallable | ||
| import numpy as np | ||
| import numpy.typing as npt | ||
| import optype as op | ||
| from numpy._typing import _ArrayLikeFloat_co, _ArrayLikeNumber_co | ||
| from scipy._typing import AnyInt, AnyReal | ||
| from .windows._windows import _Window, _WindowNeedsParams | ||
|
|
||
| __all__ = ["check_COLA", "check_NOLA", "coherence", "csd", "istft", "lombscargle", "periodogram", "spectrogram", "stft", "welch"] | ||
|
|
||
| def lombscargle(x: Untyped, y: Untyped, freqs: Untyped, precenter: bool = False, normalize: bool = False) -> Untyped: ... | ||
| _Array_f8: TypeAlias = npt.NDArray[np.float64] | ||
| _Array_f8_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]] | ||
| _ArrayFloat: TypeAlias = npt.NDArray[np.float32 | np.float64 | np.longdouble] | ||
| _ArrayComplex: TypeAlias = npt.NDArray[np.complex64 | np.complex128 | np.clongdouble] | ||
|
|
||
| _GetWindowArgument: TypeAlias = _Window | tuple[_Window | _WindowNeedsParams, Unpack[tuple[object, ...]]] | ||
| _WindowLike: TypeAlias = _GetWindowArgument | _ArrayLikeFloat_co | ||
| _Detrend: TypeAlias = Literal["literal", "constant", False] | Callable[[npt.NDArray[np.generic]], npt.NDArray[np.generic]] | ||
| _Scaling: TypeAlias = Literal["density", "spectrum"] | ||
| _LegacyScaling: TypeAlias = Literal["psd", "spectrum"] | ||
| _Average: TypeAlias = Literal["mean", "median"] | ||
| _Boundary: TypeAlias = Literal["even", "odd", "constant", "zeros"] | None | ||
|
|
||
| def lombscargle( | ||
| x: _ArrayLikeFloat_co, | ||
| y: _ArrayLikeFloat_co, | ||
| freqs: _ArrayLikeFloat_co, | ||
| precenter: op.CanBool = False, | ||
| normalize: op.CanBool = False, | ||
| ) -> _Array_f8_1d: ... | ||
| def periodogram( | ||
| x: Untyped, | ||
| fs: float = 1.0, | ||
| window: str = "boxcar", | ||
| nfft: int | None = None, | ||
| detrend: str | Literal[False] | UntypedCallable = "constant", | ||
| return_onesided: bool = True, | ||
| scaling: str = "density", | ||
| axis: int = -1, | ||
| ) -> Untyped: ... | ||
| x: _ArrayLikeNumber_co, | ||
| fs: AnyReal = 1.0, | ||
| window: _WindowLike | None = "boxcar", | ||
| nfft: AnyInt | None = None, | ||
| detrend: _Detrend = "constant", | ||
| return_onesided: op.CanBool = True, | ||
| scaling: _Scaling = "density", | ||
| axis: op.CanIndex = -1, | ||
| ) -> tuple[_Array_f8, _ArrayFloat]: ... | ||
| def welch( | ||
| x: Untyped, | ||
| fs: float = 1.0, | ||
| window: str = "hann", | ||
| nperseg: int | None = None, | ||
| noverlap: int | None = None, | ||
| nfft: int | None = None, | ||
| detrend: str | Literal[False] | UntypedCallable = "constant", | ||
| return_onesided: bool = True, | ||
| scaling: str = "density", | ||
| axis: int = -1, | ||
| average: str = "mean", | ||
| ) -> Untyped: ... | ||
| x: _ArrayLikeNumber_co, | ||
| fs: AnyReal = 1.0, | ||
| window: _WindowLike = "hann", | ||
| nperseg: AnyInt | None = None, | ||
| noverlap: AnyInt | None = None, | ||
| nfft: AnyInt | None = None, | ||
| detrend: _Detrend = "constant", | ||
| return_onesided: op.CanBool = True, | ||
| scaling: _Scaling = "density", | ||
| axis: op.CanIndex = -1, | ||
| average: _Average = "mean", | ||
| ) -> tuple[_Array_f8, _ArrayFloat]: ... | ||
| def csd( | ||
| x: Untyped, | ||
| y: Untyped, | ||
| fs: float = 1.0, | ||
| window: str = "hann", | ||
| nperseg: int | None = None, | ||
| noverlap: int | None = None, | ||
| nfft: int | None = None, | ||
| detrend: str | Literal[False] | UntypedCallable = "constant", | ||
| return_onesided: bool = True, | ||
| scaling: str = "density", | ||
| axis: int = -1, | ||
| average: str = "mean", | ||
| ) -> Untyped: ... | ||
| x: _ArrayLikeNumber_co, | ||
| y: _ArrayLikeNumber_co, | ||
| fs: AnyReal = 1.0, | ||
| window: _WindowLike = "hann", | ||
| nperseg: AnyInt | None = None, | ||
| noverlap: AnyInt | None = None, | ||
| nfft: AnyInt | None = None, | ||
| detrend: _Detrend = "constant", | ||
| return_onesided: op.CanBool = True, | ||
| scaling: _Scaling = "density", | ||
| axis: op.CanIndex = -1, | ||
| average: _Average = "mean", | ||
| ) -> tuple[_Array_f8, _ArrayComplex]: ... | ||
|
|
||
| # | ||
| @overload | ||
| # non-complex mode (positional and keyword) | ||
| def spectrogram( | ||
| x: _ArrayLikeNumber_co, | ||
| fs: AnyReal = 1.0, | ||
| window: _WindowLike = ("tukey", 0.25), | ||
| nperseg: AnyInt | None = None, | ||
| noverlap: AnyInt | None = None, | ||
| nfft: AnyInt | None = None, | ||
| detrend: _Detrend = "constant", | ||
| return_onesided: op.CanBool = True, | ||
| scaling: _Scaling = "density", | ||
| axis: op.CanIndex = -1, | ||
| mode: Literal["psd", "magnitude", "angle", "phase"] = "psd", | ||
| ) -> tuple[_Array_f8, _Array_f8, _ArrayFloat]: ... | ||
| @overload | ||
| # complex mode (positional) | ||
| def spectrogram( | ||
| x: _ArrayLikeNumber_co, | ||
| fs: AnyReal, | ||
| window: _WindowLike, | ||
| nperseg: AnyInt | None, | ||
| noverlap: AnyInt | None, | ||
| nfft: AnyInt | None, | ||
| detrend: _Detrend, | ||
| return_onesided: op.CanBool, | ||
| scaling: _Scaling, | ||
| axis: op.CanIndex, | ||
| mode: Literal["complex"], | ||
| ) -> tuple[_Array_f8, _Array_f8, _ArrayComplex]: ... | ||
| @overload | ||
| # complex mode (keyword) | ||
| def spectrogram( | ||
| x: Untyped, | ||
| fs: float = 1.0, | ||
| window: Untyped = ("tukey", 0.25), | ||
| nperseg: int | None = None, | ||
| noverlap: int | None = None, | ||
| nfft: int | None = None, | ||
| detrend: str | Literal[False] | UntypedCallable = "constant", | ||
| return_onesided: bool = True, | ||
| scaling: str = "density", | ||
| axis: int = -1, | ||
| mode: str = "psd", | ||
| ) -> Untyped: ... | ||
| def check_COLA(window: Untyped, nperseg: int, noverlap: int, tol: float = 1e-10) -> Untyped: ... | ||
| def check_NOLA(window: Untyped, nperseg: int, noverlap: int, tol: float = 1e-10) -> Untyped: ... | ||
| x: _ArrayLikeNumber_co, | ||
| fs: AnyReal = 1.0, | ||
| window: _WindowLike = ("tukey", 0.25), | ||
| nperseg: AnyInt | None = None, | ||
| noverlap: AnyInt | None = None, | ||
| nfft: AnyInt | None = None, | ||
| detrend: _Detrend = "constant", | ||
| return_onesided: op.CanBool = True, | ||
| scaling: _Scaling = "density", | ||
| axis: op.CanIndex = -1, | ||
| *, | ||
| mode: Literal["complex"], | ||
| ) -> tuple[_Array_f8, _Array_f8, _ArrayComplex]: ... | ||
|
|
||
| # | ||
| def check_COLA( | ||
| window: _WindowLike, | ||
| nperseg: AnyInt, | ||
| noverlap: AnyInt, | ||
| tol: AnyReal = 1e-10, | ||
| ) -> np.bool_: ... | ||
| def check_NOLA( | ||
| window: _WindowLike, | ||
| nperseg: AnyInt, | ||
| noverlap: AnyInt, | ||
| tol: AnyReal = 1e-10, | ||
| ) -> np.bool_: ... | ||
| def stft( | ||
| x: Untyped, | ||
| fs: float = 1.0, | ||
| window: str = "hann", | ||
| nperseg: int = 256, | ||
| noverlap: int | None = None, | ||
| nfft: int | None = None, | ||
| detrend: bool = False, | ||
| return_onesided: bool = True, | ||
| boundary: str = "zeros", | ||
| padded: bool = True, | ||
| axis: int = -1, | ||
| scaling: str = "spectrum", | ||
| ) -> Untyped: ... | ||
| x: _ArrayLikeNumber_co, | ||
| fs: AnyReal = 1.0, | ||
| window: _WindowLike = "hann", | ||
| nperseg: AnyInt = 256, | ||
| noverlap: AnyInt | None = None, | ||
| nfft: AnyInt | None = None, | ||
| detrend: _Detrend = False, | ||
| return_onesided: op.CanBool = True, | ||
| boundary: _Boundary = "zeros", | ||
| padded: op.CanBool = True, | ||
| axis: op.CanIndex = -1, | ||
| scaling: _LegacyScaling = "spectrum", | ||
| ) -> tuple[_Array_f8, _Array_f8, _ArrayComplex]: ... | ||
|
|
||
| # | ||
| @overload | ||
| # input_onesided is `True` | ||
| def istft( | ||
| Zxx: _ArrayLikeNumber_co, | ||
| fs: AnyReal = 1.0, | ||
| window: _WindowLike = "hann", | ||
| nperseg: AnyInt | None = None, | ||
| noverlap: AnyInt | None = None, | ||
| nfft: AnyInt | None = None, | ||
| input_onesided: Literal[True, 1] = True, | ||
| boundary: op.CanBool = True, | ||
| time_axis: op.CanIndex = -1, | ||
| freq_axis: op.CanIndex = -2, | ||
| scaling: _LegacyScaling = "spectrum", | ||
| ) -> tuple[_Array_f8, _ArrayFloat]: ... | ||
| @overload | ||
| # input_onesided is `False` (positional) | ||
| def istft( | ||
| Zxx: Untyped, | ||
| fs: float = 1.0, | ||
| window: str = "hann", | ||
| nperseg: int | None = None, | ||
| noverlap: int | None = None, | ||
| nfft: int | None = None, | ||
| input_onesided: bool = True, | ||
| boundary: bool = True, | ||
| time_axis: int = -1, | ||
| freq_axis: int = -2, | ||
| scaling: str = "spectrum", | ||
| ) -> Untyped: ... | ||
| Zxx: _ArrayLikeNumber_co, | ||
| fs: AnyReal, | ||
| window: _WindowLike, | ||
| nperseg: AnyInt | None, | ||
| noverlap: AnyInt | None, | ||
| nfft: AnyInt | None, | ||
| input_onesided: Literal[False, 0], | ||
| boundary: op.CanBool = True, | ||
| time_axis: op.CanIndex = -1, | ||
| freq_axis: op.CanIndex = -2, | ||
| scaling: _LegacyScaling = "spectrum", | ||
| ) -> tuple[_Array_f8, _ArrayComplex]: ... | ||
| @overload | ||
| # input_onesided is `False` (keyword) | ||
| def istft( | ||
| Zxx: _ArrayLikeNumber_co, | ||
| fs: AnyReal = 1.0, | ||
| window: _WindowLike = "hann", | ||
| nperseg: AnyInt | None = None, | ||
| noverlap: AnyInt | None = None, | ||
| nfft: AnyInt | None = None, | ||
| *, | ||
| input_onesided: Literal[False, 0], | ||
| boundary: op.CanBool = True, | ||
| time_axis: op.CanIndex = -1, | ||
| freq_axis: op.CanIndex = -2, | ||
| scaling: _LegacyScaling = "spectrum", | ||
| ) -> tuple[_Array_f8, _ArrayComplex]: ... | ||
|
|
||
| # | ||
| def coherence( | ||
| x: Untyped, | ||
| y: Untyped, | ||
| fs: float = 1.0, | ||
| window: str = "hann", | ||
| nperseg: int | None = None, | ||
| noverlap: int | None = None, | ||
| nfft: int | None = None, | ||
| detrend: str | Literal[False] | UntypedCallable = "constant", | ||
| axis: int = -1, | ||
| ) -> Untyped: ... | ||
| x: _ArrayLikeNumber_co, | ||
| y: _ArrayLikeNumber_co, | ||
| fs: AnyReal = 1.0, | ||
| window: _WindowLike = "hann", | ||
| nperseg: AnyInt | None = None, | ||
| noverlap: AnyInt | None = None, | ||
| nfft: AnyInt | None = None, | ||
| detrend: _Detrend = "constant", | ||
| axis: op.CanIndex = -1, | ||
| ) -> tuple[_Array_f8, _ArrayFloat]: ... |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,65 @@ | ||
| from typing import Literal, TypeAlias | ||
| from typing_extensions import assert_type | ||
|
|
||
| import numpy as np | ||
| import numpy.typing as npt | ||
| import optype.numpy as onpt | ||
| from scipy.signal import istft, spectrogram | ||
|
|
||
| _Array_f8: TypeAlias = npt.NDArray[np.float64] | ||
| _ArrayFloat: TypeAlias = npt.NDArray[np.float32 | np.float64 | np.longdouble] | ||
| _ArrayComplex: TypeAlias = npt.NDArray[np.complex64 | np.complex128 | np.clongdouble] | ||
|
|
||
| array_f8_1d: onpt.Array[tuple[Literal[256]], np.float64] | ||
| array_c16_1d: onpt.Array[tuple[Literal[256]], np.complex128] | ||
| spectrogram_mode_real: Literal["psd", "magnitude", "angle", "phase"] | ||
|
|
||
| # test spectrogram function overloads | ||
| assert_type(spectrogram(array_f8_1d), tuple[_Array_f8, _Array_f8, _ArrayFloat]) | ||
| assert_type(spectrogram(array_f8_1d, mode=spectrogram_mode_real), tuple[_Array_f8, _Array_f8, _ArrayFloat]) | ||
| assert_type(spectrogram(array_f8_1d, mode="complex"), tuple[_Array_f8, _Array_f8, _ArrayComplex]) | ||
| assert_type( | ||
| spectrogram(array_f8_1d, 1.0, ("tukey", 2.5), None, None, None, "constant", True, "density", -1, "complex"), | ||
| tuple[_Array_f8, _Array_f8, _ArrayComplex], | ||
| ) | ||
|
|
||
| # test isft function overloads | ||
| assert_type(istft(array_c16_1d), tuple[_Array_f8, _ArrayFloat]) | ||
| assert_type(istft(array_c16_1d, input_onesided=True), tuple[_Array_f8, _ArrayFloat]) | ||
| assert_type(istft(array_c16_1d, 1.0, "hann", 256, 128, 256, False), tuple[_Array_f8, _ArrayComplex]) | ||
| assert_type( | ||
| istft(array_c16_1d, input_onesided=False, fs=1.0, window="hann", nperseg=256, noverlap=128, nfft=256), | ||
| tuple[_Array_f8, _ArrayComplex], | ||
| ) | ||
| assert_type( | ||
| istft( | ||
| array_c16_1d, | ||
| fs=2.0, | ||
| window=("tukey", 0.25), | ||
| nperseg=256, | ||
| noverlap=128, | ||
| nfft=256, | ||
| input_onesided=True, | ||
| boundary=False, | ||
| time_axis=-1, | ||
| freq_axis=0, | ||
| scaling="spectrum", | ||
| ), | ||
| tuple[_Array_f8, _ArrayFloat], | ||
| ) | ||
| assert_type( | ||
| istft( | ||
| array_c16_1d, | ||
| fs=2.0, | ||
| window=("tukey", 0.25), | ||
| nperseg=256, | ||
| noverlap=128, | ||
| nfft=256, | ||
| input_onesided=False, | ||
| boundary=False, | ||
| time_axis=0, | ||
| freq_axis=1, | ||
| scaling="spectrum", | ||
| ), | ||
| tuple[_Array_f8, _ArrayComplex], | ||
| ) |