_polynomial_impl.pyi 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338
  1. from typing import (
  2. Any,
  3. Literal as L,
  4. NoReturn,
  5. SupportsIndex,
  6. SupportsInt,
  7. TypeAlias,
  8. TypeVar,
  9. overload,
  10. )
  11. import numpy as np
  12. from numpy import (
  13. complex128,
  14. complexfloating,
  15. float64,
  16. floating,
  17. int32,
  18. int64,
  19. object_,
  20. poly1d,
  21. signedinteger,
  22. unsignedinteger,
  23. )
  24. from numpy._typing import (
  25. ArrayLike,
  26. NDArray,
  27. _ArrayLikeBool_co,
  28. _ArrayLikeComplex_co,
  29. _ArrayLikeFloat_co,
  30. _ArrayLikeInt_co,
  31. _ArrayLikeObject_co,
  32. _ArrayLikeUInt_co,
  33. )
  34. _T = TypeVar("_T")
  35. _2Tup: TypeAlias = tuple[_T, _T]
  36. _5Tup: TypeAlias = tuple[
  37. _T,
  38. NDArray[float64],
  39. NDArray[int32],
  40. NDArray[float64],
  41. NDArray[float64],
  42. ]
  43. __all__ = [
  44. "poly",
  45. "roots",
  46. "polyint",
  47. "polyder",
  48. "polyadd",
  49. "polysub",
  50. "polymul",
  51. "polydiv",
  52. "polyval",
  53. "poly1d",
  54. "polyfit",
  55. ]
  56. def poly(seq_of_zeros: ArrayLike) -> NDArray[floating]: ...
  57. # Returns either a float or complex array depending on the input values.
  58. # See `np.linalg.eigvals`.
  59. def roots(p: ArrayLike) -> NDArray[complexfloating] | NDArray[floating]: ...
  60. @overload
  61. def polyint(
  62. p: poly1d,
  63. m: SupportsInt | SupportsIndex = 1,
  64. k: _ArrayLikeComplex_co | _ArrayLikeObject_co | None = None,
  65. ) -> poly1d: ...
  66. @overload
  67. def polyint(
  68. p: _ArrayLikeFloat_co,
  69. m: SupportsInt | SupportsIndex = 1,
  70. k: _ArrayLikeFloat_co | None = None,
  71. ) -> NDArray[floating]: ...
  72. @overload
  73. def polyint(
  74. p: _ArrayLikeComplex_co,
  75. m: SupportsInt | SupportsIndex = 1,
  76. k: _ArrayLikeComplex_co | None = None,
  77. ) -> NDArray[complexfloating]: ...
  78. @overload
  79. def polyint(
  80. p: _ArrayLikeObject_co,
  81. m: SupportsInt | SupportsIndex = 1,
  82. k: _ArrayLikeObject_co | None = None,
  83. ) -> NDArray[object_]: ...
  84. @overload
  85. def polyder(
  86. p: poly1d,
  87. m: SupportsInt | SupportsIndex = 1,
  88. ) -> poly1d: ...
  89. @overload
  90. def polyder(
  91. p: _ArrayLikeFloat_co,
  92. m: SupportsInt | SupportsIndex = 1,
  93. ) -> NDArray[floating]: ...
  94. @overload
  95. def polyder(
  96. p: _ArrayLikeComplex_co,
  97. m: SupportsInt | SupportsIndex = 1,
  98. ) -> NDArray[complexfloating]: ...
  99. @overload
  100. def polyder(
  101. p: _ArrayLikeObject_co,
  102. m: SupportsInt | SupportsIndex = 1,
  103. ) -> NDArray[object_]: ...
  104. @overload
  105. def polyfit(
  106. x: _ArrayLikeFloat_co,
  107. y: _ArrayLikeFloat_co,
  108. deg: SupportsIndex | SupportsInt,
  109. rcond: float | None = None,
  110. full: L[False] = False,
  111. w: _ArrayLikeFloat_co | None = None,
  112. cov: L[False] = False,
  113. ) -> NDArray[float64]: ...
  114. @overload
  115. def polyfit(
  116. x: _ArrayLikeComplex_co,
  117. y: _ArrayLikeComplex_co,
  118. deg: SupportsIndex | SupportsInt,
  119. rcond: float | None = None,
  120. full: L[False] = False,
  121. w: _ArrayLikeFloat_co | None = None,
  122. cov: L[False] = False,
  123. ) -> NDArray[complex128]: ...
  124. @overload
  125. def polyfit(
  126. x: _ArrayLikeFloat_co,
  127. y: _ArrayLikeFloat_co,
  128. deg: SupportsIndex | SupportsInt,
  129. rcond: float | None = None,
  130. full: L[False] = False,
  131. w: _ArrayLikeFloat_co | None = None,
  132. *,
  133. cov: L[True, "unscaled"],
  134. ) -> _2Tup[NDArray[float64]]: ...
  135. @overload
  136. def polyfit(
  137. x: _ArrayLikeComplex_co,
  138. y: _ArrayLikeComplex_co,
  139. deg: SupportsIndex | SupportsInt,
  140. rcond: float | None = None,
  141. full: L[False] = False,
  142. w: _ArrayLikeFloat_co | None = None,
  143. *,
  144. cov: L[True, "unscaled"],
  145. ) -> _2Tup[NDArray[complex128]]: ...
  146. @overload
  147. def polyfit(
  148. x: _ArrayLikeFloat_co,
  149. y: _ArrayLikeFloat_co,
  150. deg: SupportsIndex | SupportsInt,
  151. rcond: float | None,
  152. full: L[True],
  153. w: _ArrayLikeFloat_co | None = None,
  154. cov: bool | L["unscaled"] = False,
  155. ) -> _5Tup[NDArray[float64]]: ...
  156. @overload
  157. def polyfit(
  158. x: _ArrayLikeFloat_co,
  159. y: _ArrayLikeFloat_co,
  160. deg: SupportsIndex | SupportsInt,
  161. rcond: float | None = None,
  162. *,
  163. full: L[True],
  164. w: _ArrayLikeFloat_co | None = None,
  165. cov: bool | L["unscaled"] = False,
  166. ) -> _5Tup[NDArray[float64]]: ...
  167. @overload
  168. def polyfit(
  169. x: _ArrayLikeComplex_co,
  170. y: _ArrayLikeComplex_co,
  171. deg: SupportsIndex | SupportsInt,
  172. rcond: float | None,
  173. full: L[True],
  174. w: _ArrayLikeFloat_co | None = None,
  175. cov: bool | L["unscaled"] = False,
  176. ) -> _5Tup[NDArray[complex128]]: ...
  177. @overload
  178. def polyfit(
  179. x: _ArrayLikeComplex_co,
  180. y: _ArrayLikeComplex_co,
  181. deg: SupportsIndex | SupportsInt,
  182. rcond: float | None = None,
  183. *,
  184. full: L[True],
  185. w: _ArrayLikeFloat_co | None = None,
  186. cov: bool | L["unscaled"] = False,
  187. ) -> _5Tup[NDArray[complex128]]: ...
  188. @overload
  189. def polyval(
  190. p: _ArrayLikeBool_co,
  191. x: _ArrayLikeBool_co,
  192. ) -> NDArray[int64]: ...
  193. @overload
  194. def polyval(
  195. p: _ArrayLikeUInt_co,
  196. x: _ArrayLikeUInt_co,
  197. ) -> NDArray[unsignedinteger]: ...
  198. @overload
  199. def polyval(
  200. p: _ArrayLikeInt_co,
  201. x: _ArrayLikeInt_co,
  202. ) -> NDArray[signedinteger]: ...
  203. @overload
  204. def polyval(
  205. p: _ArrayLikeFloat_co,
  206. x: _ArrayLikeFloat_co,
  207. ) -> NDArray[floating]: ...
  208. @overload
  209. def polyval(
  210. p: _ArrayLikeComplex_co,
  211. x: _ArrayLikeComplex_co,
  212. ) -> NDArray[complexfloating]: ...
  213. @overload
  214. def polyval(
  215. p: _ArrayLikeObject_co,
  216. x: _ArrayLikeObject_co,
  217. ) -> NDArray[object_]: ...
  218. @overload
  219. def polyadd(
  220. a1: poly1d,
  221. a2: _ArrayLikeComplex_co | _ArrayLikeObject_co,
  222. ) -> poly1d: ...
  223. @overload
  224. def polyadd(
  225. a1: _ArrayLikeComplex_co | _ArrayLikeObject_co,
  226. a2: poly1d,
  227. ) -> poly1d: ...
  228. @overload
  229. def polyadd(
  230. a1: _ArrayLikeBool_co,
  231. a2: _ArrayLikeBool_co,
  232. ) -> NDArray[np.bool]: ...
  233. @overload
  234. def polyadd(
  235. a1: _ArrayLikeUInt_co,
  236. a2: _ArrayLikeUInt_co,
  237. ) -> NDArray[unsignedinteger]: ...
  238. @overload
  239. def polyadd(
  240. a1: _ArrayLikeInt_co,
  241. a2: _ArrayLikeInt_co,
  242. ) -> NDArray[signedinteger]: ...
  243. @overload
  244. def polyadd(
  245. a1: _ArrayLikeFloat_co,
  246. a2: _ArrayLikeFloat_co,
  247. ) -> NDArray[floating]: ...
  248. @overload
  249. def polyadd(
  250. a1: _ArrayLikeComplex_co,
  251. a2: _ArrayLikeComplex_co,
  252. ) -> NDArray[complexfloating]: ...
  253. @overload
  254. def polyadd(
  255. a1: _ArrayLikeObject_co,
  256. a2: _ArrayLikeObject_co,
  257. ) -> NDArray[object_]: ...
  258. @overload
  259. def polysub(
  260. a1: poly1d,
  261. a2: _ArrayLikeComplex_co | _ArrayLikeObject_co,
  262. ) -> poly1d: ...
  263. @overload
  264. def polysub(
  265. a1: _ArrayLikeComplex_co | _ArrayLikeObject_co,
  266. a2: poly1d,
  267. ) -> poly1d: ...
  268. @overload
  269. def polysub(
  270. a1: _ArrayLikeBool_co,
  271. a2: _ArrayLikeBool_co,
  272. ) -> NoReturn: ...
  273. @overload
  274. def polysub(
  275. a1: _ArrayLikeUInt_co,
  276. a2: _ArrayLikeUInt_co,
  277. ) -> NDArray[unsignedinteger]: ...
  278. @overload
  279. def polysub(
  280. a1: _ArrayLikeInt_co,
  281. a2: _ArrayLikeInt_co,
  282. ) -> NDArray[signedinteger]: ...
  283. @overload
  284. def polysub(
  285. a1: _ArrayLikeFloat_co,
  286. a2: _ArrayLikeFloat_co,
  287. ) -> NDArray[floating]: ...
  288. @overload
  289. def polysub(
  290. a1: _ArrayLikeComplex_co,
  291. a2: _ArrayLikeComplex_co,
  292. ) -> NDArray[complexfloating]: ...
  293. @overload
  294. def polysub(
  295. a1: _ArrayLikeObject_co,
  296. a2: _ArrayLikeObject_co,
  297. ) -> NDArray[object_]: ...
  298. # NOTE: Not an alias, but they do have the same signature (that we can reuse)
  299. polymul = polyadd
  300. @overload
  301. def polydiv(
  302. u: poly1d,
  303. v: _ArrayLikeComplex_co | _ArrayLikeObject_co,
  304. ) -> _2Tup[poly1d]: ...
  305. @overload
  306. def polydiv(
  307. u: _ArrayLikeComplex_co | _ArrayLikeObject_co,
  308. v: poly1d,
  309. ) -> _2Tup[poly1d]: ...
  310. @overload
  311. def polydiv(
  312. u: _ArrayLikeFloat_co,
  313. v: _ArrayLikeFloat_co,
  314. ) -> _2Tup[NDArray[floating]]: ...
  315. @overload
  316. def polydiv(
  317. u: _ArrayLikeComplex_co,
  318. v: _ArrayLikeComplex_co,
  319. ) -> _2Tup[NDArray[complexfloating]]: ...
  320. @overload
  321. def polydiv(
  322. u: _ArrayLikeObject_co,
  323. v: _ArrayLikeObject_co,
  324. ) -> _2Tup[NDArray[Any]]: ...