Skip to content

Commit

Permalink
🎨 signal: code simplicifactions related to window
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorenham
jorenham committed Dec 28, 2024
1 parent 1201a5e commit 6d2db1d
Show file tree
Hide file tree
Showing 3 changed files with 155 additions and 203 deletions.
65 changes: 37 additions & 28 deletions scipy-stubs/signal/_short_time_fft.pyi
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,24 @@
from collections.abc import Callable
from typing import Any, Literal, TypeAlias
from typing_extensions import Self, Unpack
from typing_extensions import Self

import numpy as np
import optype.numpy as onp
from .windows._windows import _Window, _WindowNeedsParams
from .windows._windows import _ToWindow

__all__ = ["ShortTimeFFT"]

_InexactND: TypeAlias = onp.ArrayND[np.inexact[Any]]

_ScaleTo: TypeAlias = Literal["magnitude", "psd"]
_Detr: TypeAlias = (
Literal["linear", "constant"]
| Callable[[onp.ArrayND[np.float64]], onp.ToComplexND]
| Callable[[onp.ArrayND[np.complex128]], onp.ToComplexND]
)

###

# awkward naming, but this matches the "attempts at type-aliases" in the implementation
PAD_TYPE: TypeAlias = Literal["zeros", "edge", "even", "odd"]
FFT_MODE_TYPE: TypeAlias = Literal["twosided", "centered", "onesided", "onesided2X"]
Expand All @@ -16,44 +27,44 @@ class ShortTimeFFT:
def __init__(
self,
/,
win: onp.ArrayND[np.inexact[Any]],
win: _InexactND,
hop: int,
fs: float,
*,
fft_mode: FFT_MODE_TYPE = "onesided",
mfft: int | None = None,
dual_win: onp.ArrayND[np.inexact[Any]] | None = None,
scale_to: Literal["magnitude", "psd"] | None = None,
dual_win: _InexactND | None = None,
scale_to: _ScaleTo | None = None,
phase_shift: int | None = 0,
) -> None: ...
@classmethod
def from_dual(
cls,
dual_win: onp.ArrayND[np.inexact[Any]],
dual_win: _InexactND,
hop: int,
fs: float,
*,
fft_mode: FFT_MODE_TYPE = "onesided",
mfft: int | None = None,
scale_to: Literal["magnitude", "psd"] | None = None,
scale_to: _ScaleTo | None = None,
phase_shift: int | None = 0,
) -> Self: ...
@classmethod
def from_window(
cls,
win_param: _Window | onp.ToFloat | tuple[_Window | _WindowNeedsParams, Unpack[tuple[object, ...]]],
win_param: _ToWindow,
fs: float,
nperseg: int,
noverlap: int,
*,
symmetric_win: bool = False,
fft_mode: FFT_MODE_TYPE = "onesided",
mfft: int | None = None,
scale_to: Literal["magnitude", "psd"] | None = None,
scale_to: _ScaleTo | None = None,
phase_shift: int | None = 0,
) -> Self: ...
@property
def win(self, /) -> onp.ArrayND[np.inexact[Any]]: ...
def win(self, /) -> _InexactND: ...
@property
def hop(self, /) -> int: ...
@property
Expand All @@ -73,64 +84,62 @@ class ShortTimeFFT:
@mfft.setter
def mfft(self, /, n_: int) -> None: ...
@property
def scaling(self, /) -> Literal["magnitude", "psd"] | None: ...
def scale_to(self, /, scaling: Literal["magnitude", "psd"]) -> None: ...
def scaling(self, /) -> _ScaleTo | None: ...
def scale_to(self, /, scaling: _ScaleTo) -> None: ...
@property
def phase_shift(self, /) -> int | None: ...
@phase_shift.setter
def phase_shift(self, /, v: int | None) -> None: ...
def stft(
self,
/,
x: onp.ArrayND[np.inexact[Any]],
x: _InexactND,
p0: int | None = None,
p1: int | None = None,
*,
k_offset: int = 0,
padding: PAD_TYPE = "zeros",
axis: int = -1,
) -> onp.ArrayND[np.inexact[Any]]: ...
) -> _InexactND: ...
def stft_detrend(
self,
/,
x: onp.ArrayND[np.inexact[Any]],
detr: Callable[[onp.ArrayND[np.inexact[Any]]], onp.ArrayND[np.inexact[Any]]] | Literal["linear", "constant"] | None,
x: _InexactND,
detr: _Detr | None,
p0: int | None = None,
p1: int | None = None,
*,
k_offset: int = 0,
padding: PAD_TYPE = "zeros",
axis: int = -1,
) -> onp.ArrayND[np.inexact[Any]]: ...
) -> _InexactND: ...
def spectrogram(
self,
/,
x: onp.ArrayND[np.inexact[Any]],
y: onp.ArrayND[np.inexact[Any]] | None = None,
detr: Callable[[onp.ArrayND[np.inexact[Any]]], onp.ArrayND[np.inexact[Any]]]
| Literal["linear", "constant"]
| None = None,
x: _InexactND,
y: _InexactND | None = None,
detr: _Detr | None = None,
*,
p0: int | None = None,
p1: int | None = None,
k_offset: int = 0,
padding: PAD_TYPE = "zeros",
axis: int = -1,
) -> onp.ArrayND[np.inexact[Any]]: ...
) -> _InexactND: ...
@property
def dual_win(self, /) -> onp.ArrayND[np.inexact[Any]]: ...
def dual_win(self, /) -> _InexactND: ...
@property
def invertible(self, /) -> bool: ...
def istft(
self,
/,
S: onp.ArrayND[np.inexact[Any]],
S: _InexactND,
k0: int = 0,
k1: int | None = None,
*,
f_axis: int = -2,
t_axis: int = -1,
) -> onp.ArrayND[np.inexact[Any]]: ...
) -> _InexactND: ...
@property
def fac_magnitude(self, /) -> float: ...
@property
Expand All @@ -152,7 +161,7 @@ class ShortTimeFFT:
@property
def delta_t(self, /) -> float: ...
def p_range(self, /, n: int, p0: int | None = None, p1: int | None = None) -> tuple[int, int]: ...
def t(self, /, n: int, p0: int | None = None, p1: int | None = None, k_offset: int = 0) -> onp.ArrayND[np.inexact[Any]]: ...
def t(self, /, n: int, p0: int | None = None, p1: int | None = None, k_offset: int = 0) -> _InexactND: ...
def nearest_k_p(self, /, k: int, left: bool = True) -> int: ...
@property
def delta_f(self, /) -> float: ...
Expand All @@ -161,7 +170,7 @@ class ShortTimeFFT:
@property
def onesided_fft(self, /) -> bool: ...
@property
def f(self, /) -> onp.ArrayND[np.inexact[Any]]: ...
def f(self, /) -> _InexactND: ...
def extent(
self,
/,
Expand Down
44 changes: 15 additions & 29 deletions scipy-stubs/signal/_spectral_py.pyi
View file
Original file line number Diff line number Diff line change
@@ -1,32 +1,18 @@
from collections.abc import Callable
from typing import Literal, TypeAlias, overload
from typing_extensions import Unpack

import numpy as np
import optype as op
import optype.numpy as onp
from .windows._windows import _Window, _WindowNeedsParams
from .windows._windows import _ToWindow

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

_Float1D: TypeAlias = onp.Array1D[np.float64]
_FloatND: TypeAlias = onp.ArrayND[np.float64]
_FloatingND: TypeAlias = onp.ArrayND[np.float32 | np.float64 | np.longdouble]
_CFloatingND: TypeAlias = onp.ArrayND[np.complex64 | np.complex128 | np.clongdouble]

_GetWindowArgument: TypeAlias = _Window | tuple[_Window | _WindowNeedsParams, Unpack[tuple[object, ...]]]
_WindowLike: TypeAlias = _GetWindowArgument | onp.ToFloat1D
_Detrend: TypeAlias = Literal["literal", "constant", False] | Callable[[onp.ArrayND], onp.ArrayND]
_Scaling: TypeAlias = Literal["density", "spectrum"]
_LegacyScaling: TypeAlias = Literal["psd", "spectrum"]
Expand All @@ -50,7 +36,7 @@ def lombscargle(
def periodogram(
x: onp.ToComplexND,
fs: onp.ToFloat = 1.0,
window: _WindowLike | None = "boxcar",
window: _ToWindow | None = "boxcar",
nfft: onp.ToInt | None = None,
detrend: _Detrend = "constant",
return_onesided: op.CanBool = True,
Expand All @@ -62,7 +48,7 @@ def periodogram(
def welch(
x: onp.ToComplexND,
fs: onp.ToFloat = 1.0,
window: _WindowLike = "hann",
window: _ToWindow = "hann",
nperseg: onp.ToInt | None = None,
noverlap: onp.ToInt | None = None,
nfft: onp.ToInt | None = None,
Expand All @@ -78,7 +64,7 @@ def csd(
x: onp.ToComplexND,
y: onp.ToComplexND,
fs: onp.ToFloat = 1.0,
window: _WindowLike = "hann",
window: _ToWindow = "hann",
nperseg: onp.ToInt | None = None,
noverlap: onp.ToInt | None = None,
nfft: onp.ToInt | None = None,
Expand All @@ -94,7 +80,7 @@ def csd(
def spectrogram(
x: onp.ToComplexND,
fs: onp.ToFloat = 1.0,
window: _WindowLike = ("tukey", 0.25),
window: _ToWindow = ("tukey", 0.25),
nperseg: onp.ToInt | None = None,
noverlap: onp.ToInt | None = None,
nfft: onp.ToInt | None = None,
Expand All @@ -108,7 +94,7 @@ def spectrogram(
def spectrogram(
x: onp.ToComplexND,
fs: onp.ToFloat,
window: _WindowLike,
window: _ToWindow,
nperseg: onp.ToInt | None,
noverlap: onp.ToInt | None,
nfft: onp.ToInt | None,
Expand All @@ -122,7 +108,7 @@ def spectrogram(
def spectrogram(
x: onp.ToComplexND,
fs: onp.ToFloat = 1.0,
window: _WindowLike = ("tukey", 0.25),
window: _ToWindow = ("tukey", 0.25),
nperseg: onp.ToInt | None = None,
noverlap: onp.ToInt | None = None,
nfft: onp.ToInt | None = None,
Expand All @@ -135,14 +121,14 @@ def spectrogram(
) -> tuple[_FloatND, _FloatND, _CFloatingND]: ...

#
def check_COLA(window: _WindowLike, nperseg: onp.ToInt, noverlap: onp.ToInt, tol: onp.ToFloat = 1e-10) -> np.bool_: ...
def check_NOLA(window: _WindowLike, nperseg: onp.ToInt, noverlap: onp.ToInt, tol: onp.ToFloat = 1e-10) -> np.bool_: ...
def check_COLA(window: _ToWindow, nperseg: onp.ToInt, noverlap: onp.ToInt, tol: onp.ToFloat = 1e-10) -> np.bool_: ...
def check_NOLA(window: _ToWindow, nperseg: onp.ToInt, noverlap: onp.ToInt, tol: onp.ToFloat = 1e-10) -> np.bool_: ...

#
def stft(
x: onp.ToComplexND,
fs: onp.ToFloat = 1.0,
window: _WindowLike = "hann",
window: _ToWindow = "hann",
nperseg: onp.ToInt = 256,
noverlap: onp.ToInt | None = None,
nfft: onp.ToInt | None = None,
Expand All @@ -159,7 +145,7 @@ def stft(
def istft(
Zxx: onp.ToComplexND,
fs: onp.ToFloat = 1.0,
window: _WindowLike = "hann",
window: _ToWindow = "hann",
nperseg: onp.ToInt | None = None,
noverlap: onp.ToInt | None = None,
nfft: onp.ToInt | None = None,
Expand All @@ -173,7 +159,7 @@ def istft(
def istft(
Zxx: onp.ToComplexND,
fs: onp.ToFloat,
window: _WindowLike,
window: _ToWindow,
nperseg: onp.ToInt | None,
noverlap: onp.ToInt | None,
nfft: onp.ToInt | None,
Expand All @@ -187,7 +173,7 @@ def istft(
def istft(
Zxx: onp.ToComplexND,
fs: onp.ToFloat = 1.0,
window: _WindowLike = "hann",
window: _ToWindow = "hann",
nperseg: onp.ToInt | None = None,
noverlap: onp.ToInt | None = None,
nfft: onp.ToInt | None = None,
Expand All @@ -204,7 +190,7 @@ def coherence(
x: onp.ToComplexND,
y: onp.ToComplexND,
fs: onp.ToFloat = 1.0,
window: _WindowLike = "hann",
window: _ToWindow = "hann",
nperseg: onp.ToInt | None = None,
noverlap: onp.ToInt | None = None,
nfft: onp.ToInt | None = None,
Expand Down
Loading

0 comments on commit 6d2db1d

Please sign in to comment.