_twodim_base_impl.pyi 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437
  1. from collections.abc import Callable, Sequence
  2. from typing import (
  3. Any,
  4. TypeAlias,
  5. overload,
  6. TypeVar,
  7. Literal as L,
  8. )
  9. import numpy as np
  10. from numpy import (
  11. generic,
  12. timedelta64,
  13. datetime64,
  14. int_,
  15. intp,
  16. float64,
  17. complex128,
  18. signedinteger,
  19. floating,
  20. complexfloating,
  21. object_,
  22. _OrderCF,
  23. )
  24. from numpy._typing import (
  25. DTypeLike,
  26. _DTypeLike,
  27. ArrayLike,
  28. _ArrayLike,
  29. NDArray,
  30. _SupportsArray,
  31. _SupportsArrayFunc,
  32. _ArrayLikeInt_co,
  33. _ArrayLikeFloat_co,
  34. _ArrayLikeComplex_co,
  35. _ArrayLikeObject_co,
  36. )
  37. __all__ = [
  38. "diag",
  39. "diagflat",
  40. "eye",
  41. "fliplr",
  42. "flipud",
  43. "tri",
  44. "triu",
  45. "tril",
  46. "vander",
  47. "histogram2d",
  48. "mask_indices",
  49. "tril_indices",
  50. "tril_indices_from",
  51. "triu_indices",
  52. "triu_indices_from",
  53. ]
  54. ###
  55. _T = TypeVar("_T")
  56. _SCT = TypeVar("_SCT", bound=generic)
  57. _SCT_complex = TypeVar("_SCT_complex", bound=np.complexfloating)
  58. _SCT_inexact = TypeVar("_SCT_inexact", bound=np.inexact)
  59. _SCT_number_co = TypeVar("_SCT_number_co", bound=_Number_co)
  60. # The returned arrays dtype must be compatible with `np.equal`
  61. _MaskFunc: TypeAlias = Callable[[NDArray[int_], _T], NDArray[_Number_co | timedelta64 | datetime64 | object_]]
  62. _Int_co: TypeAlias = np.integer | np.bool
  63. _Float_co: TypeAlias = np.floating | _Int_co
  64. _Number_co: TypeAlias = np.number | np.bool
  65. _ArrayLike1D: TypeAlias = _SupportsArray[np.dtype[_SCT]] | Sequence[_SCT]
  66. _ArrayLike1DInt_co: TypeAlias = _SupportsArray[np.dtype[_Int_co]] | Sequence[int | _Int_co]
  67. _ArrayLike1DFloat_co: TypeAlias = _SupportsArray[np.dtype[_Float_co]] | Sequence[float | _Float_co]
  68. _ArrayLike2DFloat_co: TypeAlias = _SupportsArray[np.dtype[_Float_co]] | Sequence[_ArrayLike1DFloat_co]
  69. _ArrayLike1DNumber_co: TypeAlias = _SupportsArray[np.dtype[_Number_co]] | Sequence[complex | _Number_co]
  70. ###
  71. @overload
  72. def fliplr(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
  73. @overload
  74. def fliplr(m: ArrayLike) -> NDArray[Any]: ...
  75. @overload
  76. def flipud(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
  77. @overload
  78. def flipud(m: ArrayLike) -> NDArray[Any]: ...
  79. @overload
  80. def eye(
  81. N: int,
  82. M: None | int = ...,
  83. k: int = ...,
  84. dtype: None = ...,
  85. order: _OrderCF = ...,
  86. *,
  87. device: None | L["cpu"] = ...,
  88. like: None | _SupportsArrayFunc = ...,
  89. ) -> NDArray[float64]: ...
  90. @overload
  91. def eye(
  92. N: int,
  93. M: None | int,
  94. k: int,
  95. dtype: _DTypeLike[_SCT],
  96. order: _OrderCF = ...,
  97. *,
  98. device: None | L["cpu"] = ...,
  99. like: None | _SupportsArrayFunc = ...,
  100. ) -> NDArray[_SCT]: ...
  101. @overload
  102. def eye(
  103. N: int,
  104. M: None | int = ...,
  105. k: int = ...,
  106. *,
  107. dtype: _DTypeLike[_SCT],
  108. order: _OrderCF = ...,
  109. device: None | L["cpu"] = ...,
  110. like: None | _SupportsArrayFunc = ...,
  111. ) -> NDArray[_SCT]: ...
  112. @overload
  113. def eye(
  114. N: int,
  115. M: None | int = ...,
  116. k: int = ...,
  117. dtype: DTypeLike = ...,
  118. order: _OrderCF = ...,
  119. *,
  120. device: None | L["cpu"] = ...,
  121. like: None | _SupportsArrayFunc = ...,
  122. ) -> NDArray[Any]: ...
  123. @overload
  124. def diag(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ...
  125. @overload
  126. def diag(v: ArrayLike, k: int = ...) -> NDArray[Any]: ...
  127. @overload
  128. def diagflat(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ...
  129. @overload
  130. def diagflat(v: ArrayLike, k: int = ...) -> NDArray[Any]: ...
  131. @overload
  132. def tri(
  133. N: int,
  134. M: None | int = ...,
  135. k: int = ...,
  136. dtype: None = ...,
  137. *,
  138. like: None | _SupportsArrayFunc = ...
  139. ) -> NDArray[float64]: ...
  140. @overload
  141. def tri(
  142. N: int,
  143. M: None | int,
  144. k: int,
  145. dtype: _DTypeLike[_SCT],
  146. *,
  147. like: None | _SupportsArrayFunc = ...
  148. ) -> NDArray[_SCT]: ...
  149. @overload
  150. def tri(
  151. N: int,
  152. M: None | int = ...,
  153. k: int = ...,
  154. *,
  155. dtype: _DTypeLike[_SCT],
  156. like: None | _SupportsArrayFunc = ...
  157. ) -> NDArray[_SCT]: ...
  158. @overload
  159. def tri(
  160. N: int,
  161. M: None | int = ...,
  162. k: int = ...,
  163. dtype: DTypeLike = ...,
  164. *,
  165. like: None | _SupportsArrayFunc = ...
  166. ) -> NDArray[Any]: ...
  167. @overload
  168. def tril(m: _ArrayLike[_SCT], k: int = 0) -> NDArray[_SCT]: ...
  169. @overload
  170. def tril(m: ArrayLike, k: int = 0) -> NDArray[Any]: ...
  171. @overload
  172. def triu(m: _ArrayLike[_SCT], k: int = 0) -> NDArray[_SCT]: ...
  173. @overload
  174. def triu(m: ArrayLike, k: int = 0) -> NDArray[Any]: ...
  175. @overload
  176. def vander( # type: ignore[misc]
  177. x: _ArrayLikeInt_co,
  178. N: None | int = ...,
  179. increasing: bool = ...,
  180. ) -> NDArray[signedinteger[Any]]: ...
  181. @overload
  182. def vander( # type: ignore[misc]
  183. x: _ArrayLikeFloat_co,
  184. N: None | int = ...,
  185. increasing: bool = ...,
  186. ) -> NDArray[floating[Any]]: ...
  187. @overload
  188. def vander(
  189. x: _ArrayLikeComplex_co,
  190. N: None | int = ...,
  191. increasing: bool = ...,
  192. ) -> NDArray[complexfloating[Any, Any]]: ...
  193. @overload
  194. def vander(
  195. x: _ArrayLikeObject_co,
  196. N: None | int = ...,
  197. increasing: bool = ...,
  198. ) -> NDArray[object_]: ...
  199. @overload
  200. def histogram2d(
  201. x: _ArrayLike1D[_SCT_complex],
  202. y: _ArrayLike1D[_SCT_complex | _Float_co],
  203. bins: int | Sequence[int] = ...,
  204. range: None | _ArrayLike2DFloat_co = ...,
  205. density: None | bool = ...,
  206. weights: None | _ArrayLike1DFloat_co = ...,
  207. ) -> tuple[
  208. NDArray[float64],
  209. NDArray[_SCT_complex],
  210. NDArray[_SCT_complex],
  211. ]: ...
  212. @overload
  213. def histogram2d(
  214. x: _ArrayLike1D[_SCT_complex | _Float_co],
  215. y: _ArrayLike1D[_SCT_complex],
  216. bins: int | Sequence[int] = ...,
  217. range: None | _ArrayLike2DFloat_co = ...,
  218. density: None | bool = ...,
  219. weights: None | _ArrayLike1DFloat_co = ...,
  220. ) -> tuple[
  221. NDArray[float64],
  222. NDArray[_SCT_complex],
  223. NDArray[_SCT_complex],
  224. ]: ...
  225. @overload
  226. def histogram2d(
  227. x: _ArrayLike1D[_SCT_inexact],
  228. y: _ArrayLike1D[_SCT_inexact | _Int_co],
  229. bins: int | Sequence[int] = ...,
  230. range: None | _ArrayLike2DFloat_co = ...,
  231. density: None | bool = ...,
  232. weights: None | _ArrayLike1DFloat_co = ...,
  233. ) -> tuple[
  234. NDArray[float64],
  235. NDArray[_SCT_inexact],
  236. NDArray[_SCT_inexact],
  237. ]: ...
  238. @overload
  239. def histogram2d(
  240. x: _ArrayLike1D[_SCT_inexact | _Int_co],
  241. y: _ArrayLike1D[_SCT_inexact],
  242. bins: int | Sequence[int] = ...,
  243. range: None | _ArrayLike2DFloat_co = ...,
  244. density: None | bool = ...,
  245. weights: None | _ArrayLike1DFloat_co = ...,
  246. ) -> tuple[
  247. NDArray[float64],
  248. NDArray[_SCT_inexact],
  249. NDArray[_SCT_inexact],
  250. ]: ...
  251. @overload
  252. def histogram2d(
  253. x: _ArrayLike1DInt_co | Sequence[float | int],
  254. y: _ArrayLike1DInt_co | Sequence[float | int],
  255. bins: int | Sequence[int] = ...,
  256. range: None | _ArrayLike2DFloat_co = ...,
  257. density: None | bool = ...,
  258. weights: None | _ArrayLike1DFloat_co = ...,
  259. ) -> tuple[
  260. NDArray[float64],
  261. NDArray[float64],
  262. NDArray[float64],
  263. ]: ...
  264. @overload
  265. def histogram2d(
  266. x: Sequence[complex | float | int],
  267. y: Sequence[complex | float | int],
  268. bins: int | Sequence[int] = ...,
  269. range: None | _ArrayLike2DFloat_co = ...,
  270. density: None | bool = ...,
  271. weights: None | _ArrayLike1DFloat_co = ...,
  272. ) -> tuple[
  273. NDArray[float64],
  274. NDArray[complex128 | float64],
  275. NDArray[complex128 | float64],
  276. ]: ...
  277. @overload
  278. def histogram2d(
  279. x: _ArrayLike1DNumber_co,
  280. y: _ArrayLike1DNumber_co,
  281. bins: _ArrayLike1D[_SCT_number_co] | Sequence[_ArrayLike1D[_SCT_number_co]],
  282. range: None | _ArrayLike2DFloat_co = ...,
  283. density: None | bool = ...,
  284. weights: None | _ArrayLike1DFloat_co = ...,
  285. ) -> tuple[
  286. NDArray[float64],
  287. NDArray[_SCT_number_co],
  288. NDArray[_SCT_number_co],
  289. ]: ...
  290. @overload
  291. def histogram2d(
  292. x: _ArrayLike1D[_SCT_inexact],
  293. y: _ArrayLike1D[_SCT_inexact],
  294. bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
  295. range: None | _ArrayLike2DFloat_co = ...,
  296. density: None | bool = ...,
  297. weights: None | _ArrayLike1DFloat_co = ...,
  298. ) -> tuple[
  299. NDArray[float64],
  300. NDArray[_SCT_number_co | _SCT_inexact],
  301. NDArray[_SCT_number_co | _SCT_inexact],
  302. ]: ...
  303. @overload
  304. def histogram2d(
  305. x: _ArrayLike1DInt_co | Sequence[float | int],
  306. y: _ArrayLike1DInt_co | Sequence[float | int],
  307. bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
  308. range: None | _ArrayLike2DFloat_co = ...,
  309. density: None | bool = ...,
  310. weights: None | _ArrayLike1DFloat_co = ...,
  311. ) -> tuple[
  312. NDArray[float64],
  313. NDArray[_SCT_number_co | float64],
  314. NDArray[_SCT_number_co | float64],
  315. ]: ...
  316. @overload
  317. def histogram2d(
  318. x: Sequence[complex | float | int],
  319. y: Sequence[complex | float | int],
  320. bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
  321. range: None | _ArrayLike2DFloat_co = ...,
  322. density: None | bool = ...,
  323. weights: None | _ArrayLike1DFloat_co = ...,
  324. ) -> tuple[
  325. NDArray[float64],
  326. NDArray[_SCT_number_co | complex128 | float64],
  327. NDArray[_SCT_number_co | complex128 | float64] ,
  328. ]: ...
  329. @overload
  330. def histogram2d(
  331. x: _ArrayLike1DNumber_co,
  332. y: _ArrayLike1DNumber_co,
  333. bins: Sequence[Sequence[bool]],
  334. range: None | _ArrayLike2DFloat_co = ...,
  335. density: None | bool = ...,
  336. weights: None | _ArrayLike1DFloat_co = ...,
  337. ) -> tuple[
  338. NDArray[float64],
  339. NDArray[np.bool],
  340. NDArray[np.bool],
  341. ]: ...
  342. @overload
  343. def histogram2d(
  344. x: _ArrayLike1DNumber_co,
  345. y: _ArrayLike1DNumber_co,
  346. bins: Sequence[Sequence[int | bool]],
  347. range: None | _ArrayLike2DFloat_co = ...,
  348. density: None | bool = ...,
  349. weights: None | _ArrayLike1DFloat_co = ...,
  350. ) -> tuple[
  351. NDArray[float64],
  352. NDArray[np.int_ | np.bool],
  353. NDArray[np.int_ | np.bool],
  354. ]: ...
  355. @overload
  356. def histogram2d(
  357. x: _ArrayLike1DNumber_co,
  358. y: _ArrayLike1DNumber_co,
  359. bins: Sequence[Sequence[float | int | bool]],
  360. range: None | _ArrayLike2DFloat_co = ...,
  361. density: None | bool = ...,
  362. weights: None | _ArrayLike1DFloat_co = ...,
  363. ) -> tuple[
  364. NDArray[float64],
  365. NDArray[np.float64 | np.int_ | np.bool],
  366. NDArray[np.float64 | np.int_ | np.bool],
  367. ]: ...
  368. @overload
  369. def histogram2d(
  370. x: _ArrayLike1DNumber_co,
  371. y: _ArrayLike1DNumber_co,
  372. bins: Sequence[Sequence[complex | float | int | bool]],
  373. range: None | _ArrayLike2DFloat_co = ...,
  374. density: None | bool = ...,
  375. weights: None | _ArrayLike1DFloat_co = ...,
  376. ) -> tuple[
  377. NDArray[float64],
  378. NDArray[np.complex128 | np.float64 | np.int_ | np.bool],
  379. NDArray[np.complex128 | np.float64 | np.int_ | np.bool],
  380. ]: ...
  381. # NOTE: we're assuming/demanding here the `mask_func` returns
  382. # an ndarray of shape `(n, n)`; otherwise there is the possibility
  383. # of the output tuple having more or less than 2 elements
  384. @overload
  385. def mask_indices(
  386. n: int,
  387. mask_func: _MaskFunc[int],
  388. k: int = ...,
  389. ) -> tuple[NDArray[intp], NDArray[intp]]: ...
  390. @overload
  391. def mask_indices(
  392. n: int,
  393. mask_func: _MaskFunc[_T],
  394. k: _T,
  395. ) -> tuple[NDArray[intp], NDArray[intp]]: ...
  396. def tril_indices(
  397. n: int,
  398. k: int = ...,
  399. m: None | int = ...,
  400. ) -> tuple[NDArray[int_], NDArray[int_]]: ...
  401. def tril_indices_from(
  402. arr: NDArray[Any],
  403. k: int = ...,
  404. ) -> tuple[NDArray[int_], NDArray[int_]]: ...
  405. def triu_indices(
  406. n: int,
  407. k: int = ...,
  408. m: None | int = ...,
  409. ) -> tuple[NDArray[int_], NDArray[int_]]: ...
  410. def triu_indices_from(
  411. arr: NDArray[Any],
  412. k: int = ...,
  413. ) -> tuple[NDArray[int_], NDArray[int_]]: ...