Skip to content

Commit

Permalink
✨ signal: complete _signaltools (#379)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorenham
jorenham authored Dec 28, 2024
2 parents 754f1d6 + 5b03c74 commit e74ac08
Show file tree
Hide file tree
Showing 6 changed files with 834 additions and 321 deletions.
4 changes: 2 additions & 2 deletions scipy-stubs/_typing.pyi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ __all__ = [
"AnyShape",
"ByteOrder",
"Casting",
"CorrelateMode",
"ConvMode",
"EnterNoneMixin",
"EnterSelfMixin",
"FileLike",
Expand Down Expand Up @@ -102,7 +102,7 @@ ByteOrder: TypeAlias = Literal["S", "<", "little", ">", "big", "=", "native", "|
OrderCF: TypeAlias = Literal["C", "F"]
OrderKACF: TypeAlias = Literal["K", "A", OrderCF]
Casting: TypeAlias = Literal["no", "equiv", "safe", "same_kind", "unsafe"]
CorrelateMode: TypeAlias = Literal["valid", "same", "full"]
ConvMode: TypeAlias = Literal["valid", "same", "full"]

# scipy literals
NanPolicy: TypeAlias = Literal["raise", "propagate", "omit"]
Expand Down
20 changes: 10 additions & 10 deletions scipy-stubs/linalg/_special_matrices.pyi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import numpy as np
import numpy.typing as npt
import optype.numpy as onp
import optype.typing as opt
from scipy._typing import CorrelateMode
from scipy._typing import ConvMode

__all__ = [
"block_diag",
Expand Down Expand Up @@ -148,23 +148,23 @@ def companion(a: _ToArrayNd[_SCT]) -> _Array2dPlus[_SCT]: ...

#
@overload
def convolution_matrix(a: _JustInt1d, n: onp.ToInt, mode: CorrelateMode = "full") -> _Int2d: ...
def convolution_matrix(a: _JustInt1d, n: onp.ToInt, mode: ConvMode = "full") -> _Int2d: ...
@overload
def convolution_matrix(a: _JustInt2dPlus, n: onp.ToInt, mode: CorrelateMode = "full") -> _Int3dPlus: ...
def convolution_matrix(a: _JustInt2dPlus, n: onp.ToInt, mode: ConvMode = "full") -> _Int3dPlus: ...
@overload
def convolution_matrix(a: _JustFloat1d, n: onp.ToInt, mode: CorrelateMode = "full") -> _Float2d: ...
def convolution_matrix(a: _JustFloat1d, n: onp.ToInt, mode: ConvMode = "full") -> _Float2d: ...
@overload
def convolution_matrix(a: _JustFloat2dPlus, n: onp.ToInt, mode: CorrelateMode = "full") -> _Float3dPlus: ...
def convolution_matrix(a: _JustFloat2dPlus, n: onp.ToInt, mode: ConvMode = "full") -> _Float3dPlus: ...
@overload
def convolution_matrix(a: _JustComplex1d, n: onp.ToInt, mode: CorrelateMode = "full") -> _Complex2d: ...
def convolution_matrix(a: _JustComplex1d, n: onp.ToInt, mode: ConvMode = "full") -> _Complex2d: ...
@overload
def convolution_matrix(a: _JustComplex2dPlus, n: onp.ToInt, mode: CorrelateMode = "full") -> _Complex3dPlus: ...
def convolution_matrix(a: _JustComplex2dPlus, n: onp.ToInt, mode: ConvMode = "full") -> _Complex3dPlus: ...
@overload
def convolution_matrix(a: _ToArrayStrict1d[_SCT], n: onp.ToInt, mode: CorrelateMode = "full") -> _Array2d[_SCT]: ...
def convolution_matrix(a: _ToArrayStrict1d[_SCT], n: onp.ToInt, mode: ConvMode = "full") -> _Array2d[_SCT]: ...
@overload
def convolution_matrix(a: _ToArrayStrict1dPlus[_SCT], n: onp.ToInt, mode: CorrelateMode = "full") -> _Array3dPlus[_SCT]: ...
def convolution_matrix(a: _ToArrayStrict1dPlus[_SCT], n: onp.ToInt, mode: ConvMode = "full") -> _Array3dPlus[_SCT]: ...
@overload
def convolution_matrix(a: _ToArrayNd[_SCT], n: onp.ToInt, mode: CorrelateMode = "full") -> _Array2dPlus[_SCT]: ...
def convolution_matrix(a: _ToArrayNd[_SCT], n: onp.ToInt, mode: ConvMode = "full") -> _Array2dPlus[_SCT]: ...

#
@overload
Expand Down
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
Loading

0 comments on commit e74ac08

Please sign in to comment.