polyutils.pyi 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431
  1. from collections.abc import Callable, Iterable, Sequence
  2. from typing import (
  3. Any,
  4. Final,
  5. Literal,
  6. SupportsIndex,
  7. TypeAlias,
  8. TypeVar,
  9. overload,
  10. )
  11. import numpy as np
  12. import numpy.typing as npt
  13. from numpy._typing import (
  14. _FloatLike_co,
  15. _NumberLike_co,
  16. _ArrayLikeFloat_co,
  17. _ArrayLikeComplex_co,
  18. )
  19. from ._polytypes import (
  20. _AnyInt,
  21. _CoefLike_co,
  22. _Array2,
  23. _Tuple2,
  24. _FloatSeries,
  25. _CoefSeries,
  26. _ComplexSeries,
  27. _ObjectSeries,
  28. _ComplexArray,
  29. _FloatArray,
  30. _CoefArray,
  31. _ObjectArray,
  32. _SeriesLikeInt_co,
  33. _SeriesLikeFloat_co,
  34. _SeriesLikeComplex_co,
  35. _SeriesLikeCoef_co,
  36. _ArrayLikeCoef_co,
  37. _FuncBinOp,
  38. _FuncValND,
  39. _FuncVanderND,
  40. )
  41. __all__: Final[Sequence[str]] = [
  42. "as_series",
  43. "format_float",
  44. "getdomain",
  45. "mapdomain",
  46. "mapparms",
  47. "trimcoef",
  48. "trimseq",
  49. ]
  50. _AnyLineF: TypeAlias = Callable[
  51. [_CoefLike_co, _CoefLike_co],
  52. _CoefArray,
  53. ]
  54. _AnyMulF: TypeAlias = Callable[
  55. [npt.ArrayLike, npt.ArrayLike],
  56. _CoefArray,
  57. ]
  58. _AnyVanderF: TypeAlias = Callable[
  59. [npt.ArrayLike, SupportsIndex],
  60. _CoefArray,
  61. ]
  62. @overload
  63. def as_series(
  64. alist: npt.NDArray[np.integer[Any]] | _FloatArray,
  65. trim: bool = ...,
  66. ) -> list[_FloatSeries]: ...
  67. @overload
  68. def as_series(
  69. alist: _ComplexArray,
  70. trim: bool = ...,
  71. ) -> list[_ComplexSeries]: ...
  72. @overload
  73. def as_series(
  74. alist: _ObjectArray,
  75. trim: bool = ...,
  76. ) -> list[_ObjectSeries]: ...
  77. @overload
  78. def as_series( # type: ignore[overload-overlap]
  79. alist: Iterable[_FloatArray | npt.NDArray[np.integer[Any]]],
  80. trim: bool = ...,
  81. ) -> list[_FloatSeries]: ...
  82. @overload
  83. def as_series(
  84. alist: Iterable[_ComplexArray],
  85. trim: bool = ...,
  86. ) -> list[_ComplexSeries]: ...
  87. @overload
  88. def as_series(
  89. alist: Iterable[_ObjectArray],
  90. trim: bool = ...,
  91. ) -> list[_ObjectSeries]: ...
  92. @overload
  93. def as_series( # type: ignore[overload-overlap]
  94. alist: Iterable[_SeriesLikeFloat_co | float],
  95. trim: bool = ...,
  96. ) -> list[_FloatSeries]: ...
  97. @overload
  98. def as_series(
  99. alist: Iterable[_SeriesLikeComplex_co | complex],
  100. trim: bool = ...,
  101. ) -> list[_ComplexSeries]: ...
  102. @overload
  103. def as_series(
  104. alist: Iterable[_SeriesLikeCoef_co | object],
  105. trim: bool = ...,
  106. ) -> list[_ObjectSeries]: ...
  107. _T_seq = TypeVar("_T_seq", bound=_CoefArray | Sequence[_CoefLike_co])
  108. def trimseq(seq: _T_seq) -> _T_seq: ...
  109. @overload
  110. def trimcoef( # type: ignore[overload-overlap]
  111. c: npt.NDArray[np.integer[Any]] | _FloatArray,
  112. tol: _FloatLike_co = ...,
  113. ) -> _FloatSeries: ...
  114. @overload
  115. def trimcoef(
  116. c: _ComplexArray,
  117. tol: _FloatLike_co = ...,
  118. ) -> _ComplexSeries: ...
  119. @overload
  120. def trimcoef(
  121. c: _ObjectArray,
  122. tol: _FloatLike_co = ...,
  123. ) -> _ObjectSeries: ...
  124. @overload
  125. def trimcoef( # type: ignore[overload-overlap]
  126. c: _SeriesLikeFloat_co | float,
  127. tol: _FloatLike_co = ...,
  128. ) -> _FloatSeries: ...
  129. @overload
  130. def trimcoef(
  131. c: _SeriesLikeComplex_co | complex,
  132. tol: _FloatLike_co = ...,
  133. ) -> _ComplexSeries: ...
  134. @overload
  135. def trimcoef(
  136. c: _SeriesLikeCoef_co | object,
  137. tol: _FloatLike_co = ...,
  138. ) -> _ObjectSeries: ...
  139. @overload
  140. def getdomain( # type: ignore[overload-overlap]
  141. x: _FloatArray | npt.NDArray[np.integer[Any]],
  142. ) -> _Array2[np.float64]: ...
  143. @overload
  144. def getdomain(
  145. x: _ComplexArray,
  146. ) -> _Array2[np.complex128]: ...
  147. @overload
  148. def getdomain(
  149. x: _ObjectArray,
  150. ) -> _Array2[np.object_]: ...
  151. @overload
  152. def getdomain( # type: ignore[overload-overlap]
  153. x: _SeriesLikeFloat_co | float,
  154. ) -> _Array2[np.float64]: ...
  155. @overload
  156. def getdomain(
  157. x: _SeriesLikeComplex_co | complex,
  158. ) -> _Array2[np.complex128]: ...
  159. @overload
  160. def getdomain(
  161. x: _SeriesLikeCoef_co | object,
  162. ) -> _Array2[np.object_]: ...
  163. @overload
  164. def mapparms( # type: ignore[overload-overlap]
  165. old: npt.NDArray[np.floating[Any] | np.integer[Any]],
  166. new: npt.NDArray[np.floating[Any] | np.integer[Any]],
  167. ) -> _Tuple2[np.floating[Any]]: ...
  168. @overload
  169. def mapparms(
  170. old: npt.NDArray[np.number[Any]],
  171. new: npt.NDArray[np.number[Any]],
  172. ) -> _Tuple2[np.complexfloating[Any, Any]]: ...
  173. @overload
  174. def mapparms(
  175. old: npt.NDArray[np.object_ | np.number[Any]],
  176. new: npt.NDArray[np.object_ | np.number[Any]],
  177. ) -> _Tuple2[object]: ...
  178. @overload
  179. def mapparms( # type: ignore[overload-overlap]
  180. old: Sequence[float],
  181. new: Sequence[float],
  182. ) -> _Tuple2[float]: ...
  183. @overload
  184. def mapparms(
  185. old: Sequence[complex],
  186. new: Sequence[complex],
  187. ) -> _Tuple2[complex]: ...
  188. @overload
  189. def mapparms(
  190. old: _SeriesLikeFloat_co,
  191. new: _SeriesLikeFloat_co,
  192. ) -> _Tuple2[np.floating[Any]]: ...
  193. @overload
  194. def mapparms(
  195. old: _SeriesLikeComplex_co,
  196. new: _SeriesLikeComplex_co,
  197. ) -> _Tuple2[np.complexfloating[Any, Any]]: ...
  198. @overload
  199. def mapparms(
  200. old: _SeriesLikeCoef_co,
  201. new: _SeriesLikeCoef_co,
  202. ) -> _Tuple2[object]: ...
  203. @overload
  204. def mapdomain( # type: ignore[overload-overlap]
  205. x: _FloatLike_co,
  206. old: _SeriesLikeFloat_co,
  207. new: _SeriesLikeFloat_co,
  208. ) -> np.floating[Any]: ...
  209. @overload
  210. def mapdomain(
  211. x: _NumberLike_co,
  212. old: _SeriesLikeComplex_co,
  213. new: _SeriesLikeComplex_co,
  214. ) -> np.complexfloating[Any, Any]: ...
  215. @overload
  216. def mapdomain( # type: ignore[overload-overlap]
  217. x: npt.NDArray[np.floating[Any] | np.integer[Any]],
  218. old: npt.NDArray[np.floating[Any] | np.integer[Any]],
  219. new: npt.NDArray[np.floating[Any] | np.integer[Any]],
  220. ) -> _FloatSeries: ...
  221. @overload
  222. def mapdomain(
  223. x: npt.NDArray[np.number[Any]],
  224. old: npt.NDArray[np.number[Any]],
  225. new: npt.NDArray[np.number[Any]],
  226. ) -> _ComplexSeries: ...
  227. @overload
  228. def mapdomain(
  229. x: npt.NDArray[np.object_ | np.number[Any]],
  230. old: npt.NDArray[np.object_ | np.number[Any]],
  231. new: npt.NDArray[np.object_ | np.number[Any]],
  232. ) -> _ObjectSeries: ...
  233. @overload
  234. def mapdomain( # type: ignore[overload-overlap]
  235. x: _SeriesLikeFloat_co,
  236. old: _SeriesLikeFloat_co,
  237. new: _SeriesLikeFloat_co,
  238. ) -> _FloatSeries: ...
  239. @overload
  240. def mapdomain(
  241. x: _SeriesLikeComplex_co,
  242. old: _SeriesLikeComplex_co,
  243. new: _SeriesLikeComplex_co,
  244. ) -> _ComplexSeries: ...
  245. @overload
  246. def mapdomain(
  247. x: _SeriesLikeCoef_co,
  248. old:_SeriesLikeCoef_co,
  249. new: _SeriesLikeCoef_co,
  250. ) -> _ObjectSeries: ...
  251. @overload
  252. def mapdomain(
  253. x: _CoefLike_co,
  254. old: _SeriesLikeCoef_co,
  255. new: _SeriesLikeCoef_co,
  256. ) -> object: ...
  257. def _nth_slice(
  258. i: SupportsIndex,
  259. ndim: SupportsIndex,
  260. ) -> tuple[None | slice, ...]: ...
  261. _vander_nd: _FuncVanderND[Literal["_vander_nd"]]
  262. _vander_nd_flat: _FuncVanderND[Literal["_vander_nd_flat"]]
  263. # keep in sync with `._polytypes._FuncFromRoots`
  264. @overload
  265. def _fromroots( # type: ignore[overload-overlap]
  266. line_f: _AnyLineF,
  267. mul_f: _AnyMulF,
  268. roots: _SeriesLikeFloat_co,
  269. ) -> _FloatSeries: ...
  270. @overload
  271. def _fromroots(
  272. line_f: _AnyLineF,
  273. mul_f: _AnyMulF,
  274. roots: _SeriesLikeComplex_co,
  275. ) -> _ComplexSeries: ...
  276. @overload
  277. def _fromroots(
  278. line_f: _AnyLineF,
  279. mul_f: _AnyMulF,
  280. roots: _SeriesLikeCoef_co,
  281. ) -> _ObjectSeries: ...
  282. @overload
  283. def _fromroots(
  284. line_f: _AnyLineF,
  285. mul_f: _AnyMulF,
  286. roots: _SeriesLikeCoef_co,
  287. ) -> _CoefSeries: ...
  288. _valnd: _FuncValND[Literal["_valnd"]]
  289. _gridnd: _FuncValND[Literal["_gridnd"]]
  290. # keep in sync with `_polytypes._FuncBinOp`
  291. @overload
  292. def _div( # type: ignore[overload-overlap]
  293. mul_f: _AnyMulF,
  294. c1: _SeriesLikeFloat_co,
  295. c2: _SeriesLikeFloat_co,
  296. ) -> _Tuple2[_FloatSeries]: ...
  297. @overload
  298. def _div(
  299. mul_f: _AnyMulF,
  300. c1: _SeriesLikeComplex_co,
  301. c2: _SeriesLikeComplex_co,
  302. ) -> _Tuple2[_ComplexSeries]: ...
  303. @overload
  304. def _div(
  305. mul_f: _AnyMulF,
  306. c1: _SeriesLikeCoef_co,
  307. c2: _SeriesLikeCoef_co,
  308. ) -> _Tuple2[_ObjectSeries]: ...
  309. @overload
  310. def _div(
  311. mul_f: _AnyMulF,
  312. c1: _SeriesLikeCoef_co,
  313. c2: _SeriesLikeCoef_co,
  314. ) -> _Tuple2[_CoefSeries]: ...
  315. _add: Final[_FuncBinOp]
  316. _sub: Final[_FuncBinOp]
  317. # keep in sync with `_polytypes._FuncPow`
  318. @overload
  319. def _pow( # type: ignore[overload-overlap]
  320. mul_f: _AnyMulF,
  321. c: _SeriesLikeFloat_co,
  322. pow: _AnyInt,
  323. maxpower: None | _AnyInt = ...,
  324. ) -> _FloatSeries: ...
  325. @overload
  326. def _pow(
  327. mul_f: _AnyMulF,
  328. c: _SeriesLikeComplex_co,
  329. pow: _AnyInt,
  330. maxpower: None | _AnyInt = ...,
  331. ) -> _ComplexSeries: ...
  332. @overload
  333. def _pow(
  334. mul_f: _AnyMulF,
  335. c: _SeriesLikeCoef_co,
  336. pow: _AnyInt,
  337. maxpower: None | _AnyInt = ...,
  338. ) -> _ObjectSeries: ...
  339. @overload
  340. def _pow(
  341. mul_f: _AnyMulF,
  342. c: _SeriesLikeCoef_co,
  343. pow: _AnyInt,
  344. maxpower: None | _AnyInt = ...,
  345. ) -> _CoefSeries: ...
  346. # keep in sync with `_polytypes._FuncFit`
  347. @overload
  348. def _fit( # type: ignore[overload-overlap]
  349. vander_f: _AnyVanderF,
  350. x: _SeriesLikeFloat_co,
  351. y: _ArrayLikeFloat_co,
  352. deg: _SeriesLikeInt_co,
  353. domain: None | _SeriesLikeFloat_co = ...,
  354. rcond: None | _FloatLike_co = ...,
  355. full: Literal[False] = ...,
  356. w: None | _SeriesLikeFloat_co = ...,
  357. ) -> _FloatArray: ...
  358. @overload
  359. def _fit(
  360. vander_f: _AnyVanderF,
  361. x: _SeriesLikeComplex_co,
  362. y: _ArrayLikeComplex_co,
  363. deg: _SeriesLikeInt_co,
  364. domain: None | _SeriesLikeComplex_co = ...,
  365. rcond: None | _FloatLike_co = ...,
  366. full: Literal[False] = ...,
  367. w: None | _SeriesLikeComplex_co = ...,
  368. ) -> _ComplexArray: ...
  369. @overload
  370. def _fit(
  371. vander_f: _AnyVanderF,
  372. x: _SeriesLikeCoef_co,
  373. y: _ArrayLikeCoef_co,
  374. deg: _SeriesLikeInt_co,
  375. domain: None | _SeriesLikeCoef_co = ...,
  376. rcond: None | _FloatLike_co = ...,
  377. full: Literal[False] = ...,
  378. w: None | _SeriesLikeCoef_co = ...,
  379. ) -> _CoefArray: ...
  380. @overload
  381. def _fit(
  382. vander_f: _AnyVanderF,
  383. x: _SeriesLikeCoef_co,
  384. y: _SeriesLikeCoef_co,
  385. deg: _SeriesLikeInt_co,
  386. domain: None | _SeriesLikeCoef_co,
  387. rcond: None | _FloatLike_co ,
  388. full: Literal[True],
  389. /,
  390. w: None | _SeriesLikeCoef_co = ...,
  391. ) -> tuple[_CoefSeries, Sequence[np.inexact[Any] | np.int32]]: ...
  392. @overload
  393. def _fit(
  394. vander_f: _AnyVanderF,
  395. x: _SeriesLikeCoef_co,
  396. y: _SeriesLikeCoef_co,
  397. deg: _SeriesLikeInt_co,
  398. domain: None | _SeriesLikeCoef_co = ...,
  399. rcond: None | _FloatLike_co = ...,
  400. *,
  401. full: Literal[True],
  402. w: None | _SeriesLikeCoef_co = ...,
  403. ) -> tuple[_CoefSeries, Sequence[np.inexact[Any] | np.int32]]: ...
  404. def _as_int(x: SupportsIndex, desc: str) -> int: ...
  405. def format_float(x: _FloatLike_co, parens: bool = ...) -> str: ...