fix and improve scipy.signal._peak_finding

scipy-stubs/signal/_peak_finding.pyi

@@ -1,10 +1,11 @@
-from collections.abc import Callable
-from typing import Any, Concatenate, Literal, TypeAlias, TypedDict, TypeVar
+from collections.abc import Callable, Sequence
+from typing import Any, Concatenate, Literal, TypeAlias, TypedDict, TypeVar, type_check_only
 
 import numpy as np
 import numpy.typing as npt
 import optype as op
-from numpy._typing import _ArrayLikeFloat_co, _ArrayLikeInt_co
+import optype.numpy as onpt
+from numpy._typing import _ArrayLikeBool_co, _ArrayLikeFloat_co, _ArrayLikeInt_co, _ArrayLikeNumber_co
 from scipy._typing import AnyInt, AnyReal
 
 __all__ = ["argrelextrema", "argrelmax", "argrelmin", "find_peaks", "find_peaks_cwt", "peak_prominences", "peak_widths"]
@@ -14,73 +15,76 @@ _SCT = TypeVar("_SCT", bound=np.generic)
 _Array_n: TypeAlias = npt.NDArray[np.intp]
 _Array_n_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.intp]]
 _Array_f8: TypeAlias = npt.NDArray[np.float64]
+_Array_f8_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]]
 _Mode: TypeAlias = Literal["clip", "wrap"]
 
-_ProminencesResult: TypeAlias = tuple[_Array_f8, _Array_n, _Array_n]
-_WidthsResult: TypeAlias = tuple[_Array_f8, _Array_f8, _Array_f8, _Array_f8]
-_WaveletOutput: TypeAlias = op.CanGetitem[slice, npt.ArrayLike]
-_WaveletFunction: TypeAlias = Callable[Concatenate[int | np.int_ | np.float64, AnyReal, ...], _WaveletOutput]
+_ArgRel: TypeAlias = tuple[_Array_n, ...]
+_PeakProminences: TypeAlias = tuple[_Array_f8, _Array_n, _Array_n]
+_PeakWidths: TypeAlias = tuple[_Array_f8, _Array_f8, _Array_f8, _Array_f8]
+# TODO(jorenham): Narrow down the parameter types (because contravariant)
+_WaveletFunc: TypeAlias = (
+    Callable[Concatenate[int, float, ...], _ArrayLikeNumber_co]
+    | Callable[Concatenate[np.intp, np.float64, ...], _ArrayLikeNumber_co]
+)
 
+@type_check_only
 class _FindPeaksResultsDict(TypedDict, total=False):
-    peak_heights: _Array_f8
-    left_thresholds: _Array_f8
-    right_thresholds: _Array_f8
-    prominences: _Array_f8
-    left_bases: _Array_n
-    right_bases: _Array_n
-    width_heights: _Array_f8
-    left_ips: _Array_f8
-    right_ips: _Array_f8
-    plateau_sizes: _Array_n
-    left_edges: _Array_n
-    right_edges: _Array_n
+    plateau_sizes: _Array_n_1d
+    left_edges: _Array_n_1d
+    right_edges: _Array_n_1d
 
-def argrelmin(
-    data: npt.NDArray[np.generic],
-    axis: op.CanIndex = 0,
-    order: AnyInt = 1,
-    mode: _Mode = "clip",
-) -> tuple[_Array_n, ...]: ...
-def argrelmax(
-    data: npt.NDArray[np.generic],
-    axis: op.CanIndex = 0,
-    order: AnyInt = 1,
-    mode: _Mode = "clip",
-) -> tuple[_Array_n, ...]: ...
+    peak_heights: _Array_f8_1d
+
+    left_thresholds: _Array_f8_1d
+    right_thresholds: _Array_f8_1d
+
+    prominences: _Array_f8_1d
+    left_bases: _Array_n_1d
+    right_bases: _Array_n_1d
+
+    widths: _Array_f8_1d
+    width_heights: _Array_f8_1d
+    left_ips: _Array_f8_1d
+    right_ips: _Array_f8_1d
+
+###
+
+def argrelmin(data: onpt.Array, axis: op.CanIndex = 0, order: AnyInt = 1, mode: _Mode = "clip") -> _ArgRel: ...
+def argrelmax(data: onpt.Array, axis: op.CanIndex = 0, order: AnyInt = 1, mode: _Mode = "clip") -> _ArgRel: ...
 def argrelextrema(
-    data: npt.NDArray[np.generic],
-    comparator: Callable[[npt.NDArray[_SCT], npt.NDArray[_SCT]], npt.NDArray[np.bool_]],
+    data: onpt.Array,
+    comparator: Callable[[npt.NDArray[_SCT], npt.NDArray[_SCT]], _ArrayLikeBool_co],
     axis: op.CanIndex = 0,
     order: AnyInt = 1,
     mode: _Mode = "clip",
-) -> tuple[_Array_n, ...]: ...
-def peak_prominences(
-    x: npt.ArrayLike,
-    peaks: _ArrayLikeInt_co,
-    wlen: AnyReal | None = None,
-) -> _ProminencesResult: ...
+) -> _ArgRel: ...
+
+#
+def peak_prominences(x: npt.ArrayLike, peaks: _ArrayLikeInt_co, wlen: AnyReal | None = None) -> _PeakProminences: ...
 def peak_widths(
     x: npt.ArrayLike,
     peaks: _ArrayLikeInt_co,
     rel_height: AnyReal = 0.5,
-    prominence_data: _ProminencesResult | None = None,
+    prominence_data: _PeakProminences | None = None,
     wlen: AnyReal | None = None,
-) -> _WidthsResult: ...
+) -> _PeakWidths: ...
+
+#
 def find_peaks(
     x: npt.ArrayLike,
-    height: _ArrayLikeFloat_co | tuple[AnyReal | None, AnyReal | None] | None = None,
-    threshold: _ArrayLikeFloat_co | tuple[AnyReal | None, AnyReal | None] | None = None,
+    height: _ArrayLikeFloat_co | Sequence[AnyReal | None] | None = None,
+    threshold: _ArrayLikeFloat_co | Sequence[AnyReal | None] | None = None,
     distance: AnyReal | None = None,
-    prominence: _ArrayLikeFloat_co | tuple[AnyReal | None, AnyReal | None] | None = None,
-    width: _ArrayLikeFloat_co | tuple[AnyReal | None, AnyReal | None] | None = None,
+    prominence: _ArrayLikeFloat_co | Sequence[AnyReal | None] | None = None,
+    width: _ArrayLikeFloat_co | Sequence[AnyReal | None] | None = None,
     wlen: AnyReal | None = None,
     rel_height: AnyReal = 0.5,
-    plateau_size: _ArrayLikeInt_co | tuple[AnyInt | None, AnyInt | None] | None = None,
-) -> tuple[_Array_n, _FindPeaksResultsDict]: ...
+    plateau_size: _ArrayLikeInt_co | Sequence[AnyInt | None] | None = None,
+) -> tuple[_Array_n_1d, _FindPeaksResultsDict]: ...
 def find_peaks_cwt(
-    vector: npt.NDArray[np.generic],
+    vector: onpt.Array,
     widths: _ArrayLikeFloat_co,
-    wavelet: _WaveletFunction | None = None,
+    wavelet: _WaveletFunc | None = None,
     max_distances: npt.NDArray[np.floating[Any] | np.integer[Any]] | None = None,
     gap_thresh: AnyReal | None = None,
     min_length: AnyInt | None = None,

scipy-stubs/signal/_peak_finding_utils.pyi

@@ -0,0 +1,25 @@
+# scipy/signal/_peak_finding_utils.pyx
+from typing import TypeAlias
+
+import numpy as np
+import optype.numpy as onpt
+
+__all__ = ["_local_maxima_1d", "_peak_prominences", "_peak_widths", "_select_by_peak_distance"]
+
+_Array_b_1d: TypeAlias = onpt.Array[tuple[int], np.bool_]
+_Array_n_1d: TypeAlias = onpt.Array[tuple[int], np.intp]
+_Array_f8_1d: TypeAlias = onpt.Array[tuple[int], np.float64]
+
+class PeakPropertyWarning(RuntimeWarning): ...  # undocumented
+
+def _local_maxima_1d(x: _Array_f8_1d) -> tuple[_Array_n_1d, _Array_n_1d, _Array_n_1d]: ...
+def _select_by_peak_distance(peaks: _Array_n_1d, priority: _Array_f8_1d, distance: np.float64) -> _Array_b_1d: ...
+def _peak_prominences(x: _Array_f8_1d, peaks: _Array_n_1d, wlen: np.intp) -> tuple[_Array_f8_1d, _Array_n_1d, _Array_n_1d]: ...
+def _peak_widths(
+    x: _Array_f8_1d,
+    peaks: _Array_n_1d,
+    rel_height: np.float64,
+    prominences: _Array_f8_1d,
+    left_bases: _Array_n_1d,
+    right_bases: _Array_n_1d,
+) -> tuple[_Array_f8_1d, _Array_f8_1d, _Array_f8_1d, _Array_f8_1d]: ...