Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

signal: Add type stubs to _spectral_py.pyi. #157

Merged
merged 18 commits into from
Nov 4, 2024
Merged

signal: Add type stubs to _spectral_py.pyi. #157

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
b32f27d
`signal`: Add type stubs for `check_NOLA` and `check_COLA`.
pavyamsiri Nov 4, 2024
f7ba439
`signal`: Add type stubs for function `csd`.
pavyamsiri Nov 4, 2024
58cfd75
`signal`: Add type stubs for `lombscargle`, `periodogram` and `welch`.
pavyamsiri Nov 4, 2024
9f24afa
`signal`: Add type stubs for `spectrogram`.
pavyamsiri Nov 4, 2024
3c30e05
tests: Add type test for `spectrogram` from the `signal` module.
pavyamsiri Nov 4, 2024
22dad83
`signal`: Add type stubs for `stft` and `istft`.
pavyamsiri Nov 4, 2024
973820e
tests: Add type tests for function overloads of `istft`.
pavyamsiri Nov 4, 2024
bcf3d42
`signal`: Add type stubs for `coherence`.
pavyamsiri Nov 4, 2024
cc8f55a
`signal`: Clean up `_spectral_py.pyi`.
pavyamsiri Nov 4, 2024
f302e55
`signal`: Adjust type aliases in `_spectral_py.pyi`.
pavyamsiri Nov 4, 2024
d8e6561
`signal`: Tighten the return array dtypes in `_spectral_py.pyi`.
pavyamsiri Nov 4, 2024
face996
tests: Clean up `test_spectral.pyi`.
pavyamsiri Nov 4, 2024
2f7e1c6
`signal`: Fix up comment inconsistency in `_spectral_py.pyi`.
pavyamsiri Nov 4, 2024
4dfb44a
`signal`: Use the more portable names for dtypes in `_spectral_py.pyi`.
pavyamsiri Nov 4, 2024
6728216
`signal`: Use `np.bool_` instead of `np.bool`.
pavyamsiri Nov 4, 2024
847504c
`signal`: Use `np.bool_` instead of `np.bool`
pavyamsiri Nov 4, 2024
3bb8780
`signal`: Use a type alias for `_GetWindowArgument | _ArrayLikeFloat_co`
pavyamsiri Nov 4, 2024
fb5c85d
tests: Update `test_spectral.pyi` to use the correct dtypes
pavyamsiri Nov 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
282 changes: 196 additions & 86 deletions scipy-stubs/signal/_spectral_py.pyi
View file
Open in desktop
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"]
jorenham marked this conversation as resolved.
Show resolved Hide resolved
_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]: ...
65 changes: 65 additions & 0 deletions tests/signal/test_spectral.pyi
View file
Open in desktop
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],
)