polyutils.pyi 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307
  1. from collections.abc import Callable, Iterable, Sequence
  2. from typing import (
  3. Final,
  4. Literal,
  5. Protocol,
  6. SupportsIndex,
  7. TypeAlias,
  8. TypeVar,
  9. overload,
  10. type_check_only,
  11. )
  12. import numpy as np
  13. import numpy.typing as npt
  14. from numpy._typing import (
  15. _ArrayLikeComplex_co,
  16. _ArrayLikeFloat_co,
  17. _ArrayLikeObject_co,
  18. _FloatLike_co,
  19. _NumberLike_co,
  20. )
  21. from ._polytypes import (
  22. _AnyInt,
  23. _Array2,
  24. _ArrayLikeCoef_co,
  25. _CoefArray,
  26. _CoefLike_co,
  27. _CoefSeries,
  28. _ComplexArray,
  29. _ComplexSeries,
  30. _FloatArray,
  31. _FloatSeries,
  32. _FuncBinOp,
  33. _ObjectArray,
  34. _ObjectSeries,
  35. _SeriesLikeCoef_co,
  36. _SeriesLikeComplex_co,
  37. _SeriesLikeFloat_co,
  38. _SeriesLikeInt_co,
  39. _SeriesLikeObject_co,
  40. _Tuple2,
  41. )
  42. __all__ = ["as_series", "format_float", "getdomain", "mapdomain", "mapparms", "trimcoef", "trimseq"]
  43. _T = TypeVar("_T")
  44. _SeqT = TypeVar("_SeqT", bound=_CoefArray | Sequence[_CoefLike_co])
  45. _AnyLineF: TypeAlias = Callable[[float, float], _CoefArray]
  46. _AnyMulF: TypeAlias = Callable[[np.ndarray | list[int], np.ndarray], _CoefArray]
  47. _AnyVanderF: TypeAlias = Callable[[np.ndarray, int], _CoefArray]
  48. @type_check_only
  49. class _ValFunc(Protocol[_T]):
  50. def __call__(self, x: np.ndarray, c: _T, /, *, tensor: bool = True) -> _T: ...
  51. ###
  52. @overload
  53. def as_series(alist: npt.NDArray[np.integer] | _FloatArray, trim: bool = True) -> list[_FloatSeries]: ...
  54. @overload
  55. def as_series(alist: _ComplexArray, trim: bool = True) -> list[_ComplexSeries]: ...
  56. @overload
  57. def as_series(alist: _ObjectArray, trim: bool = True) -> list[_ObjectSeries]: ...
  58. @overload
  59. def as_series(alist: Iterable[_FloatArray | npt.NDArray[np.integer]], trim: bool = True) -> list[_FloatSeries]: ...
  60. @overload
  61. def as_series(alist: Iterable[_ComplexArray], trim: bool = True) -> list[_ComplexSeries]: ...
  62. @overload
  63. def as_series(alist: Iterable[_ObjectArray], trim: bool = True) -> list[_ObjectSeries]: ...
  64. @overload
  65. def as_series(alist: Iterable[_SeriesLikeFloat_co | float], trim: bool = True) -> list[_FloatSeries]: ...
  66. @overload
  67. def as_series(alist: Iterable[_SeriesLikeComplex_co | complex], trim: bool = True) -> list[_ComplexSeries]: ...
  68. @overload
  69. def as_series(alist: Iterable[_SeriesLikeCoef_co | object], trim: bool = True) -> list[_ObjectSeries]: ...
  70. #
  71. def trimseq(seq: _SeqT) -> _SeqT: ...
  72. #
  73. @overload
  74. def trimcoef(c: npt.NDArray[np.integer] | _FloatArray, tol: _FloatLike_co = 0) -> _FloatSeries: ...
  75. @overload
  76. def trimcoef(c: _ComplexArray, tol: _FloatLike_co = 0) -> _ComplexSeries: ...
  77. @overload
  78. def trimcoef(c: _ObjectArray, tol: _FloatLike_co = 0) -> _ObjectSeries: ...
  79. @overload
  80. def trimcoef(c: _SeriesLikeFloat_co | float, tol: _FloatLike_co = 0) -> _FloatSeries: ...
  81. @overload
  82. def trimcoef(c: _SeriesLikeComplex_co | complex, tol: _FloatLike_co = 0) -> _ComplexSeries: ...
  83. @overload
  84. def trimcoef(c: _SeriesLikeCoef_co | object, tol: _FloatLike_co = 0) -> _ObjectSeries: ...
  85. #
  86. @overload
  87. def getdomain(x: _FloatArray | npt.NDArray[np.integer]) -> _Array2[np.float64]: ...
  88. @overload
  89. def getdomain(x: _ComplexArray) -> _Array2[np.complex128]: ...
  90. @overload
  91. def getdomain(x: _ObjectArray) -> _Array2[np.object_]: ...
  92. @overload
  93. def getdomain(x: _SeriesLikeFloat_co | float) -> _Array2[np.float64]: ...
  94. @overload
  95. def getdomain(x: _SeriesLikeComplex_co | complex) -> _Array2[np.complex128]: ...
  96. @overload
  97. def getdomain(x: _SeriesLikeCoef_co | object) -> _Array2[np.object_]: ...
  98. #
  99. @overload
  100. def mapparms(old: npt.NDArray[np.floating | np.integer], new: npt.NDArray[np.floating | np.integer]) -> _Tuple2[np.floating]: ...
  101. @overload
  102. def mapparms(old: npt.NDArray[np.number], new: npt.NDArray[np.number]) -> _Tuple2[np.complexfloating]: ...
  103. @overload
  104. def mapparms(old: npt.NDArray[np.object_ | np.number], new: npt.NDArray[np.object_ | np.number]) -> _Tuple2[object]: ...
  105. @overload
  106. def mapparms(old: Sequence[float], new: Sequence[float]) -> _Tuple2[float]: ...
  107. @overload
  108. def mapparms(old: Sequence[complex], new: Sequence[complex]) -> _Tuple2[complex]: ...
  109. @overload
  110. def mapparms(old: _SeriesLikeFloat_co, new: _SeriesLikeFloat_co) -> _Tuple2[np.floating]: ...
  111. @overload
  112. def mapparms(old: _SeriesLikeComplex_co, new: _SeriesLikeComplex_co) -> _Tuple2[np.complexfloating]: ...
  113. @overload
  114. def mapparms(old: _SeriesLikeCoef_co, new: _SeriesLikeCoef_co) -> _Tuple2[object]: ...
  115. #
  116. @overload
  117. def mapdomain(x: _FloatLike_co, old: _SeriesLikeFloat_co, new: _SeriesLikeFloat_co) -> np.floating: ...
  118. @overload
  119. def mapdomain(x: _NumberLike_co, old: _SeriesLikeComplex_co, new: _SeriesLikeComplex_co) -> np.complexfloating: ...
  120. @overload
  121. def mapdomain(
  122. x: npt.NDArray[np.floating | np.integer],
  123. old: npt.NDArray[np.floating | np.integer],
  124. new: npt.NDArray[np.floating | np.integer],
  125. ) -> _FloatSeries: ...
  126. @overload
  127. def mapdomain(x: npt.NDArray[np.number], old: npt.NDArray[np.number], new: npt.NDArray[np.number]) -> _ComplexSeries: ...
  128. @overload
  129. def mapdomain(
  130. x: npt.NDArray[np.object_ | np.number],
  131. old: npt.NDArray[np.object_ | np.number],
  132. new: npt.NDArray[np.object_ | np.number],
  133. ) -> _ObjectSeries: ...
  134. @overload
  135. def mapdomain(x: _SeriesLikeFloat_co, old: _SeriesLikeFloat_co, new: _SeriesLikeFloat_co) -> _FloatSeries: ...
  136. @overload
  137. def mapdomain(x: _SeriesLikeComplex_co, old: _SeriesLikeComplex_co, new: _SeriesLikeComplex_co) -> _ComplexSeries: ...
  138. @overload
  139. def mapdomain(x: _SeriesLikeCoef_co, old: _SeriesLikeCoef_co, new: _SeriesLikeCoef_co) -> _ObjectSeries: ...
  140. @overload
  141. def mapdomain(x: _CoefLike_co, old: _SeriesLikeCoef_co, new: _SeriesLikeCoef_co) -> object: ...
  142. #
  143. def _nth_slice(i: SupportsIndex, ndim: SupportsIndex) -> tuple[slice | None, ...]: ...
  144. # keep in sync with `vander_nd_flat`
  145. @overload
  146. def _vander_nd(
  147. vander_fs: Sequence[_AnyVanderF],
  148. points: Sequence[_ArrayLikeFloat_co],
  149. degrees: Sequence[SupportsIndex],
  150. ) -> _FloatArray: ...
  151. @overload
  152. def _vander_nd(
  153. vander_fs: Sequence[_AnyVanderF],
  154. points: Sequence[_ArrayLikeComplex_co],
  155. degrees: Sequence[SupportsIndex],
  156. ) -> _ComplexArray: ...
  157. @overload
  158. def _vander_nd(
  159. vander_fs: Sequence[_AnyVanderF],
  160. points: Sequence[_ArrayLikeObject_co | _ArrayLikeComplex_co],
  161. degrees: Sequence[SupportsIndex],
  162. ) -> _ObjectArray: ...
  163. @overload
  164. def _vander_nd(
  165. vander_fs: Sequence[_AnyVanderF],
  166. points: Sequence[npt.ArrayLike],
  167. degrees: Sequence[SupportsIndex],
  168. ) -> _CoefArray: ...
  169. # keep in sync with `vander_nd`
  170. @overload
  171. def _vander_nd_flat(
  172. vander_fs: Sequence[_AnyVanderF],
  173. points: Sequence[_ArrayLikeFloat_co],
  174. degrees: Sequence[SupportsIndex],
  175. ) -> _FloatArray: ...
  176. @overload
  177. def _vander_nd_flat(
  178. vander_fs: Sequence[_AnyVanderF],
  179. points: Sequence[_ArrayLikeComplex_co],
  180. degrees: Sequence[SupportsIndex],
  181. ) -> _ComplexArray: ...
  182. @overload
  183. def _vander_nd_flat(
  184. vander_fs: Sequence[_AnyVanderF],
  185. points: Sequence[_ArrayLikeObject_co | _ArrayLikeComplex_co],
  186. degrees: Sequence[SupportsIndex],
  187. ) -> _ObjectArray: ...
  188. @overload
  189. def _vander_nd_flat(
  190. vander_fs: Sequence[_AnyVanderF],
  191. points: Sequence[npt.ArrayLike],
  192. degrees: Sequence[SupportsIndex],
  193. ) -> _CoefArray: ...
  194. # keep in sync with `._polytypes._FuncFromRoots`
  195. @overload
  196. def _fromroots(line_f: _AnyLineF, mul_f: _AnyMulF, roots: _SeriesLikeFloat_co) -> _FloatSeries: ...
  197. @overload
  198. def _fromroots(line_f: _AnyLineF, mul_f: _AnyMulF, roots: _SeriesLikeComplex_co) -> _ComplexSeries: ...
  199. @overload
  200. def _fromroots(line_f: _AnyLineF, mul_f: _AnyMulF, roots: _SeriesLikeObject_co) -> _ObjectSeries: ...
  201. @overload
  202. def _fromroots(line_f: _AnyLineF, mul_f: _AnyMulF, roots: _SeriesLikeCoef_co) -> _CoefSeries: ...
  203. # keep in sync with `_gridnd`
  204. def _valnd(val_f: _ValFunc[_T], c: _T, *args: npt.ArrayLike) -> _T: ...
  205. # keep in sync with `_valnd`
  206. def _gridnd(val_f: _ValFunc[_T], c: _T, *args: npt.ArrayLike) -> _T: ...
  207. # keep in sync with `_polytypes._FuncBinOp`
  208. @overload
  209. def _div(mul_f: _AnyMulF, c1: _SeriesLikeFloat_co, c2: _SeriesLikeFloat_co) -> _Tuple2[_FloatSeries]: ...
  210. @overload
  211. def _div(mul_f: _AnyMulF, c1: _SeriesLikeComplex_co, c2: _SeriesLikeComplex_co) -> _Tuple2[_ComplexSeries]: ...
  212. @overload
  213. def _div(mul_f: _AnyMulF, c1: _SeriesLikeObject_co, c2: _SeriesLikeObject_co) -> _Tuple2[_ObjectSeries]: ...
  214. @overload
  215. def _div(mul_f: _AnyMulF, c1: _SeriesLikeCoef_co, c2: _SeriesLikeCoef_co) -> _Tuple2[_CoefSeries]: ...
  216. _add: Final[_FuncBinOp] = ...
  217. _sub: Final[_FuncBinOp] = ...
  218. # keep in sync with `_polytypes._FuncPow`
  219. @overload
  220. def _pow(mul_f: _AnyMulF, c: _SeriesLikeFloat_co, pow: _AnyInt, maxpower: _AnyInt | None) -> _FloatSeries: ...
  221. @overload
  222. def _pow(mul_f: _AnyMulF, c: _SeriesLikeComplex_co, pow: _AnyInt, maxpower: _AnyInt | None) -> _ComplexSeries: ...
  223. @overload
  224. def _pow(mul_f: _AnyMulF, c: _SeriesLikeObject_co, pow: _AnyInt, maxpower: _AnyInt | None) -> _ObjectSeries: ...
  225. @overload
  226. def _pow(mul_f: _AnyMulF, c: _SeriesLikeCoef_co, pow: _AnyInt, maxpower: _AnyInt | None) -> _CoefSeries: ...
  227. # keep in sync with `_polytypes._FuncFit`
  228. @overload
  229. def _fit(
  230. vander_f: _AnyVanderF,
  231. x: _SeriesLikeFloat_co,
  232. y: _ArrayLikeFloat_co,
  233. deg: _SeriesLikeInt_co,
  234. rcond: _FloatLike_co | None = None,
  235. full: Literal[False] = False,
  236. w: _SeriesLikeFloat_co | None = None,
  237. ) -> _FloatArray: ...
  238. @overload
  239. def _fit(
  240. vander_f: _AnyVanderF,
  241. x: _SeriesLikeComplex_co,
  242. y: _ArrayLikeComplex_co,
  243. deg: _SeriesLikeInt_co,
  244. rcond: _FloatLike_co | None = None,
  245. full: Literal[False] = False,
  246. w: _SeriesLikeComplex_co | None = None,
  247. ) -> _ComplexArray: ...
  248. @overload
  249. def _fit(
  250. vander_f: _AnyVanderF,
  251. x: _SeriesLikeCoef_co,
  252. y: _ArrayLikeCoef_co,
  253. deg: _SeriesLikeInt_co,
  254. rcond: _FloatLike_co | None = None,
  255. full: Literal[False] = False,
  256. w: _SeriesLikeCoef_co | None = None,
  257. ) -> _CoefArray: ...
  258. @overload
  259. def _fit(
  260. vander_f: _AnyVanderF,
  261. x: _SeriesLikeCoef_co,
  262. y: _SeriesLikeCoef_co,
  263. deg: _SeriesLikeInt_co,
  264. rcond: _FloatLike_co | None,
  265. full: Literal[True],
  266. w: _SeriesLikeCoef_co | None = None,
  267. ) -> tuple[_CoefSeries, Sequence[np.inexact | np.int32]]: ...
  268. @overload
  269. def _fit(
  270. vander_f: _AnyVanderF,
  271. x: _SeriesLikeCoef_co,
  272. y: _SeriesLikeCoef_co,
  273. deg: _SeriesLikeInt_co,
  274. rcond: _FloatLike_co | None = None,
  275. *,
  276. full: Literal[True],
  277. w: _SeriesLikeCoef_co | None = None,
  278. ) -> tuple[_CoefSeries, Sequence[np.inexact | np.int32]]: ...
  279. #
  280. def _as_int(x: SupportsIndex, desc: str) -> int: ...
  281. #
  282. def format_float(x: _FloatLike_co, parens: bool = False) -> str: ...