polyutils.pyi 10.0 KB

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