_polynomial_impl.pyi 7.0 KB

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