_twodim_base_impl.py 33 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201
  1. """ Basic functions for manipulating 2d arrays
  2. """
  3. import functools
  4. import operator
  5. from numpy._core import iinfo, overrides
  6. from numpy._core._multiarray_umath import _array_converter
  7. from numpy._core.numeric import (
  8. arange,
  9. asanyarray,
  10. asarray,
  11. diagonal,
  12. empty,
  13. greater_equal,
  14. indices,
  15. int8,
  16. int16,
  17. int32,
  18. int64,
  19. intp,
  20. multiply,
  21. nonzero,
  22. ones,
  23. promote_types,
  24. where,
  25. zeros,
  26. )
  27. from numpy._core.overrides import finalize_array_function_like, set_module
  28. from numpy.lib._stride_tricks_impl import broadcast_to
  29. __all__ = [
  30. 'diag', 'diagflat', 'eye', 'fliplr', 'flipud', 'tri', 'triu',
  31. 'tril', 'vander', 'histogram2d', 'mask_indices', 'tril_indices',
  32. 'tril_indices_from', 'triu_indices', 'triu_indices_from', ]
  33. array_function_dispatch = functools.partial(
  34. overrides.array_function_dispatch, module='numpy')
  35. i1 = iinfo(int8)
  36. i2 = iinfo(int16)
  37. i4 = iinfo(int32)
  38. def _min_int(low, high):
  39. """ get small int that fits the range """
  40. if high <= i1.max and low >= i1.min:
  41. return int8
  42. if high <= i2.max and low >= i2.min:
  43. return int16
  44. if high <= i4.max and low >= i4.min:
  45. return int32
  46. return int64
  47. def _flip_dispatcher(m):
  48. return (m,)
  49. @array_function_dispatch(_flip_dispatcher)
  50. def fliplr(m):
  51. """
  52. Reverse the order of elements along axis 1 (left/right).
  53. For a 2-D array, this flips the entries in each row in the left/right
  54. direction. Columns are preserved, but appear in a different order than
  55. before.
  56. Parameters
  57. ----------
  58. m : array_like
  59. Input array, must be at least 2-D.
  60. Returns
  61. -------
  62. f : ndarray
  63. A view of `m` with the columns reversed. Since a view
  64. is returned, this operation is :math:`\\mathcal O(1)`.
  65. See Also
  66. --------
  67. flipud : Flip array in the up/down direction.
  68. flip : Flip array in one or more dimensions.
  69. rot90 : Rotate array counterclockwise.
  70. Notes
  71. -----
  72. Equivalent to ``m[:,::-1]`` or ``np.flip(m, axis=1)``.
  73. Requires the array to be at least 2-D.
  74. Examples
  75. --------
  76. >>> import numpy as np
  77. >>> A = np.diag([1.,2.,3.])
  78. >>> A
  79. array([[1., 0., 0.],
  80. [0., 2., 0.],
  81. [0., 0., 3.]])
  82. >>> np.fliplr(A)
  83. array([[0., 0., 1.],
  84. [0., 2., 0.],
  85. [3., 0., 0.]])
  86. >>> rng = np.random.default_rng()
  87. >>> A = rng.normal(size=(2,3,5))
  88. >>> np.all(np.fliplr(A) == A[:,::-1,...])
  89. True
  90. """
  91. m = asanyarray(m)
  92. if m.ndim < 2:
  93. raise ValueError("Input must be >= 2-d.")
  94. return m[:, ::-1]
  95. @array_function_dispatch(_flip_dispatcher)
  96. def flipud(m):
  97. """
  98. Reverse the order of elements along axis 0 (up/down).
  99. For a 2-D array, this flips the entries in each column in the up/down
  100. direction. Rows are preserved, but appear in a different order than before.
  101. Parameters
  102. ----------
  103. m : array_like
  104. Input array.
  105. Returns
  106. -------
  107. out : array_like
  108. A view of `m` with the rows reversed. Since a view is
  109. returned, this operation is :math:`\\mathcal O(1)`.
  110. See Also
  111. --------
  112. fliplr : Flip array in the left/right direction.
  113. flip : Flip array in one or more dimensions.
  114. rot90 : Rotate array counterclockwise.
  115. Notes
  116. -----
  117. Equivalent to ``m[::-1, ...]`` or ``np.flip(m, axis=0)``.
  118. Requires the array to be at least 1-D.
  119. Examples
  120. --------
  121. >>> import numpy as np
  122. >>> A = np.diag([1.0, 2, 3])
  123. >>> A
  124. array([[1., 0., 0.],
  125. [0., 2., 0.],
  126. [0., 0., 3.]])
  127. >>> np.flipud(A)
  128. array([[0., 0., 3.],
  129. [0., 2., 0.],
  130. [1., 0., 0.]])
  131. >>> rng = np.random.default_rng()
  132. >>> A = rng.normal(size=(2,3,5))
  133. >>> np.all(np.flipud(A) == A[::-1,...])
  134. True
  135. >>> np.flipud([1,2])
  136. array([2, 1])
  137. """
  138. m = asanyarray(m)
  139. if m.ndim < 1:
  140. raise ValueError("Input must be >= 1-d.")
  141. return m[::-1, ...]
  142. @finalize_array_function_like
  143. @set_module('numpy')
  144. def eye(N, M=None, k=0, dtype=float, order='C', *, device=None, like=None):
  145. """
  146. Return a 2-D array with ones on the diagonal and zeros elsewhere.
  147. Parameters
  148. ----------
  149. N : int
  150. Number of rows in the output.
  151. M : int, optional
  152. Number of columns in the output. If None, defaults to `N`.
  153. k : int, optional
  154. Index of the diagonal: 0 (the default) refers to the main diagonal,
  155. a positive value refers to an upper diagonal, and a negative value
  156. to a lower diagonal.
  157. dtype : data-type, optional
  158. Data-type of the returned array.
  159. order : {'C', 'F'}, optional
  160. Whether the output should be stored in row-major (C-style) or
  161. column-major (Fortran-style) order in memory.
  162. device : str, optional
  163. The device on which to place the created array. Default: None.
  164. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  165. .. versionadded:: 2.0.0
  166. ${ARRAY_FUNCTION_LIKE}
  167. .. versionadded:: 1.20.0
  168. Returns
  169. -------
  170. I : ndarray of shape (N,M)
  171. An array where all elements are equal to zero, except for the `k`-th
  172. diagonal, whose values are equal to one.
  173. See Also
  174. --------
  175. identity : (almost) equivalent function
  176. diag : diagonal 2-D array from a 1-D array specified by the user.
  177. Examples
  178. --------
  179. >>> import numpy as np
  180. >>> np.eye(2, dtype=int)
  181. array([[1, 0],
  182. [0, 1]])
  183. >>> np.eye(3, k=1)
  184. array([[0., 1., 0.],
  185. [0., 0., 1.],
  186. [0., 0., 0.]])
  187. """
  188. if like is not None:
  189. return _eye_with_like(
  190. like, N, M=M, k=k, dtype=dtype, order=order, device=device
  191. )
  192. if M is None:
  193. M = N
  194. m = zeros((N, M), dtype=dtype, order=order, device=device)
  195. if k >= M:
  196. return m
  197. # Ensure M and k are integers, so we don't get any surprise casting
  198. # results in the expressions `M-k` and `M+1` used below. This avoids
  199. # a problem with inputs with type (for example) np.uint64.
  200. M = operator.index(M)
  201. k = operator.index(k)
  202. if k >= 0:
  203. i = k
  204. else:
  205. i = (-k) * M
  206. m[:M - k].flat[i::M + 1] = 1
  207. return m
  208. _eye_with_like = array_function_dispatch()(eye)
  209. def _diag_dispatcher(v, k=None):
  210. return (v,)
  211. @array_function_dispatch(_diag_dispatcher)
  212. def diag(v, k=0):
  213. """
  214. Extract a diagonal or construct a diagonal array.
  215. See the more detailed documentation for ``numpy.diagonal`` if you use this
  216. function to extract a diagonal and wish to write to the resulting array;
  217. whether it returns a copy or a view depends on what version of numpy you
  218. are using.
  219. Parameters
  220. ----------
  221. v : array_like
  222. If `v` is a 2-D array, return a copy of its `k`-th diagonal.
  223. If `v` is a 1-D array, return a 2-D array with `v` on the `k`-th
  224. diagonal.
  225. k : int, optional
  226. Diagonal in question. The default is 0. Use `k>0` for diagonals
  227. above the main diagonal, and `k<0` for diagonals below the main
  228. diagonal.
  229. Returns
  230. -------
  231. out : ndarray
  232. The extracted diagonal or constructed diagonal array.
  233. See Also
  234. --------
  235. diagonal : Return specified diagonals.
  236. diagflat : Create a 2-D array with the flattened input as a diagonal.
  237. trace : Sum along diagonals.
  238. triu : Upper triangle of an array.
  239. tril : Lower triangle of an array.
  240. Examples
  241. --------
  242. >>> import numpy as np
  243. >>> x = np.arange(9).reshape((3,3))
  244. >>> x
  245. array([[0, 1, 2],
  246. [3, 4, 5],
  247. [6, 7, 8]])
  248. >>> np.diag(x)
  249. array([0, 4, 8])
  250. >>> np.diag(x, k=1)
  251. array([1, 5])
  252. >>> np.diag(x, k=-1)
  253. array([3, 7])
  254. >>> np.diag(np.diag(x))
  255. array([[0, 0, 0],
  256. [0, 4, 0],
  257. [0, 0, 8]])
  258. """
  259. v = asanyarray(v)
  260. s = v.shape
  261. if len(s) == 1:
  262. n = s[0] + abs(k)
  263. res = zeros((n, n), v.dtype)
  264. if k >= 0:
  265. i = k
  266. else:
  267. i = (-k) * n
  268. res[:n - k].flat[i::n + 1] = v
  269. return res
  270. elif len(s) == 2:
  271. return diagonal(v, k)
  272. else:
  273. raise ValueError("Input must be 1- or 2-d.")
  274. @array_function_dispatch(_diag_dispatcher)
  275. def diagflat(v, k=0):
  276. """
  277. Create a two-dimensional array with the flattened input as a diagonal.
  278. Parameters
  279. ----------
  280. v : array_like
  281. Input data, which is flattened and set as the `k`-th
  282. diagonal of the output.
  283. k : int, optional
  284. Diagonal to set; 0, the default, corresponds to the "main" diagonal,
  285. a positive (negative) `k` giving the number of the diagonal above
  286. (below) the main.
  287. Returns
  288. -------
  289. out : ndarray
  290. The 2-D output array.
  291. See Also
  292. --------
  293. diag : MATLAB work-alike for 1-D and 2-D arrays.
  294. diagonal : Return specified diagonals.
  295. trace : Sum along diagonals.
  296. Examples
  297. --------
  298. >>> import numpy as np
  299. >>> np.diagflat([[1,2], [3,4]])
  300. array([[1, 0, 0, 0],
  301. [0, 2, 0, 0],
  302. [0, 0, 3, 0],
  303. [0, 0, 0, 4]])
  304. >>> np.diagflat([1,2], 1)
  305. array([[0, 1, 0],
  306. [0, 0, 2],
  307. [0, 0, 0]])
  308. """
  309. conv = _array_converter(v)
  310. v, = conv.as_arrays(subok=False)
  311. v = v.ravel()
  312. s = len(v)
  313. n = s + abs(k)
  314. res = zeros((n, n), v.dtype)
  315. if (k >= 0):
  316. i = arange(0, n - k, dtype=intp)
  317. fi = i + k + i * n
  318. else:
  319. i = arange(0, n + k, dtype=intp)
  320. fi = i + (i - k) * n
  321. res.flat[fi] = v
  322. return conv.wrap(res)
  323. @finalize_array_function_like
  324. @set_module('numpy')
  325. def tri(N, M=None, k=0, dtype=float, *, like=None):
  326. """
  327. An array with ones at and below the given diagonal and zeros elsewhere.
  328. Parameters
  329. ----------
  330. N : int
  331. Number of rows in the array.
  332. M : int, optional
  333. Number of columns in the array.
  334. By default, `M` is taken equal to `N`.
  335. k : int, optional
  336. The sub-diagonal at and below which the array is filled.
  337. `k` = 0 is the main diagonal, while `k` < 0 is below it,
  338. and `k` > 0 is above. The default is 0.
  339. dtype : dtype, optional
  340. Data type of the returned array. The default is float.
  341. ${ARRAY_FUNCTION_LIKE}
  342. .. versionadded:: 1.20.0
  343. Returns
  344. -------
  345. tri : ndarray of shape (N, M)
  346. Array with its lower triangle filled with ones and zero elsewhere;
  347. in other words ``T[i,j] == 1`` for ``j <= i + k``, 0 otherwise.
  348. Examples
  349. --------
  350. >>> import numpy as np
  351. >>> np.tri(3, 5, 2, dtype=int)
  352. array([[1, 1, 1, 0, 0],
  353. [1, 1, 1, 1, 0],
  354. [1, 1, 1, 1, 1]])
  355. >>> np.tri(3, 5, -1)
  356. array([[0., 0., 0., 0., 0.],
  357. [1., 0., 0., 0., 0.],
  358. [1., 1., 0., 0., 0.]])
  359. """
  360. if like is not None:
  361. return _tri_with_like(like, N, M=M, k=k, dtype=dtype)
  362. if M is None:
  363. M = N
  364. m = greater_equal.outer(arange(N, dtype=_min_int(0, N)),
  365. arange(-k, M - k, dtype=_min_int(-k, M - k)))
  366. # Avoid making a copy if the requested type is already bool
  367. m = m.astype(dtype, copy=False)
  368. return m
  369. _tri_with_like = array_function_dispatch()(tri)
  370. def _trilu_dispatcher(m, k=None):
  371. return (m,)
  372. @array_function_dispatch(_trilu_dispatcher)
  373. def tril(m, k=0):
  374. """
  375. Lower triangle of an array.
  376. Return a copy of an array with elements above the `k`-th diagonal zeroed.
  377. For arrays with ``ndim`` exceeding 2, `tril` will apply to the final two
  378. axes.
  379. Parameters
  380. ----------
  381. m : array_like, shape (..., M, N)
  382. Input array.
  383. k : int, optional
  384. Diagonal above which to zero elements. `k = 0` (the default) is the
  385. main diagonal, `k < 0` is below it and `k > 0` is above.
  386. Returns
  387. -------
  388. tril : ndarray, shape (..., M, N)
  389. Lower triangle of `m`, of same shape and data-type as `m`.
  390. See Also
  391. --------
  392. triu : same thing, only for the upper triangle
  393. Examples
  394. --------
  395. >>> import numpy as np
  396. >>> np.tril([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1)
  397. array([[ 0, 0, 0],
  398. [ 4, 0, 0],
  399. [ 7, 8, 0],
  400. [10, 11, 12]])
  401. >>> np.tril(np.arange(3*4*5).reshape(3, 4, 5))
  402. array([[[ 0, 0, 0, 0, 0],
  403. [ 5, 6, 0, 0, 0],
  404. [10, 11, 12, 0, 0],
  405. [15, 16, 17, 18, 0]],
  406. [[20, 0, 0, 0, 0],
  407. [25, 26, 0, 0, 0],
  408. [30, 31, 32, 0, 0],
  409. [35, 36, 37, 38, 0]],
  410. [[40, 0, 0, 0, 0],
  411. [45, 46, 0, 0, 0],
  412. [50, 51, 52, 0, 0],
  413. [55, 56, 57, 58, 0]]])
  414. """
  415. m = asanyarray(m)
  416. mask = tri(*m.shape[-2:], k=k, dtype=bool)
  417. return where(mask, m, zeros(1, m.dtype))
  418. @array_function_dispatch(_trilu_dispatcher)
  419. def triu(m, k=0):
  420. """
  421. Upper triangle of an array.
  422. Return a copy of an array with the elements below the `k`-th diagonal
  423. zeroed. For arrays with ``ndim`` exceeding 2, `triu` will apply to the
  424. final two axes.
  425. Please refer to the documentation for `tril` for further details.
  426. See Also
  427. --------
  428. tril : lower triangle of an array
  429. Examples
  430. --------
  431. >>> import numpy as np
  432. >>> np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1)
  433. array([[ 1, 2, 3],
  434. [ 4, 5, 6],
  435. [ 0, 8, 9],
  436. [ 0, 0, 12]])
  437. >>> np.triu(np.arange(3*4*5).reshape(3, 4, 5))
  438. array([[[ 0, 1, 2, 3, 4],
  439. [ 0, 6, 7, 8, 9],
  440. [ 0, 0, 12, 13, 14],
  441. [ 0, 0, 0, 18, 19]],
  442. [[20, 21, 22, 23, 24],
  443. [ 0, 26, 27, 28, 29],
  444. [ 0, 0, 32, 33, 34],
  445. [ 0, 0, 0, 38, 39]],
  446. [[40, 41, 42, 43, 44],
  447. [ 0, 46, 47, 48, 49],
  448. [ 0, 0, 52, 53, 54],
  449. [ 0, 0, 0, 58, 59]]])
  450. """
  451. m = asanyarray(m)
  452. mask = tri(*m.shape[-2:], k=k - 1, dtype=bool)
  453. return where(mask, zeros(1, m.dtype), m)
  454. def _vander_dispatcher(x, N=None, increasing=None):
  455. return (x,)
  456. # Originally borrowed from John Hunter and matplotlib
  457. @array_function_dispatch(_vander_dispatcher)
  458. def vander(x, N=None, increasing=False):
  459. """
  460. Generate a Vandermonde matrix.
  461. The columns of the output matrix are powers of the input vector. The
  462. order of the powers is determined by the `increasing` boolean argument.
  463. Specifically, when `increasing` is False, the `i`-th output column is
  464. the input vector raised element-wise to the power of ``N - i - 1``. Such
  465. a matrix with a geometric progression in each row is named for Alexandre-
  466. Theophile Vandermonde.
  467. Parameters
  468. ----------
  469. x : array_like
  470. 1-D input array.
  471. N : int, optional
  472. Number of columns in the output. If `N` is not specified, a square
  473. array is returned (``N = len(x)``).
  474. increasing : bool, optional
  475. Order of the powers of the columns. If True, the powers increase
  476. from left to right, if False (the default) they are reversed.
  477. Returns
  478. -------
  479. out : ndarray
  480. Vandermonde matrix. If `increasing` is False, the first column is
  481. ``x^(N-1)``, the second ``x^(N-2)`` and so forth. If `increasing` is
  482. True, the columns are ``x^0, x^1, ..., x^(N-1)``.
  483. See Also
  484. --------
  485. polynomial.polynomial.polyvander
  486. Examples
  487. --------
  488. >>> import numpy as np
  489. >>> x = np.array([1, 2, 3, 5])
  490. >>> N = 3
  491. >>> np.vander(x, N)
  492. array([[ 1, 1, 1],
  493. [ 4, 2, 1],
  494. [ 9, 3, 1],
  495. [25, 5, 1]])
  496. >>> np.column_stack([x**(N-1-i) for i in range(N)])
  497. array([[ 1, 1, 1],
  498. [ 4, 2, 1],
  499. [ 9, 3, 1],
  500. [25, 5, 1]])
  501. >>> x = np.array([1, 2, 3, 5])
  502. >>> np.vander(x)
  503. array([[ 1, 1, 1, 1],
  504. [ 8, 4, 2, 1],
  505. [ 27, 9, 3, 1],
  506. [125, 25, 5, 1]])
  507. >>> np.vander(x, increasing=True)
  508. array([[ 1, 1, 1, 1],
  509. [ 1, 2, 4, 8],
  510. [ 1, 3, 9, 27],
  511. [ 1, 5, 25, 125]])
  512. The determinant of a square Vandermonde matrix is the product
  513. of the differences between the values of the input vector:
  514. >>> np.linalg.det(np.vander(x))
  515. 48.000000000000043 # may vary
  516. >>> (5-3)*(5-2)*(5-1)*(3-2)*(3-1)*(2-1)
  517. 48
  518. """
  519. x = asarray(x)
  520. if x.ndim != 1:
  521. raise ValueError("x must be a one-dimensional array or sequence.")
  522. if N is None:
  523. N = len(x)
  524. v = empty((len(x), N), dtype=promote_types(x.dtype, int))
  525. tmp = v[:, ::-1] if not increasing else v
  526. if N > 0:
  527. tmp[:, 0] = 1
  528. if N > 1:
  529. tmp[:, 1:] = x[:, None]
  530. multiply.accumulate(tmp[:, 1:], out=tmp[:, 1:], axis=1)
  531. return v
  532. def _histogram2d_dispatcher(x, y, bins=None, range=None, density=None,
  533. weights=None):
  534. yield x
  535. yield y
  536. # This terrible logic is adapted from the checks in histogram2d
  537. try:
  538. N = len(bins)
  539. except TypeError:
  540. N = 1
  541. if N == 2:
  542. yield from bins # bins=[x, y]
  543. else:
  544. yield bins
  545. yield weights
  546. @array_function_dispatch(_histogram2d_dispatcher)
  547. def histogram2d(x, y, bins=10, range=None, density=None, weights=None):
  548. """
  549. Compute the bi-dimensional histogram of two data samples.
  550. Parameters
  551. ----------
  552. x : array_like, shape (N,)
  553. An array containing the x coordinates of the points to be
  554. histogrammed.
  555. y : array_like, shape (N,)
  556. An array containing the y coordinates of the points to be
  557. histogrammed.
  558. bins : int or array_like or [int, int] or [array, array], optional
  559. The bin specification:
  560. * If int, the number of bins for the two dimensions (nx=ny=bins).
  561. * If array_like, the bin edges for the two dimensions
  562. (x_edges=y_edges=bins).
  563. * If [int, int], the number of bins in each dimension
  564. (nx, ny = bins).
  565. * If [array, array], the bin edges in each dimension
  566. (x_edges, y_edges = bins).
  567. * A combination [int, array] or [array, int], where int
  568. is the number of bins and array is the bin edges.
  569. range : array_like, shape(2,2), optional
  570. The leftmost and rightmost edges of the bins along each dimension
  571. (if not specified explicitly in the `bins` parameters):
  572. ``[[xmin, xmax], [ymin, ymax]]``. All values outside of this range
  573. will be considered outliers and not tallied in the histogram.
  574. density : bool, optional
  575. If False, the default, returns the number of samples in each bin.
  576. If True, returns the probability *density* function at the bin,
  577. ``bin_count / sample_count / bin_area``.
  578. weights : array_like, shape(N,), optional
  579. An array of values ``w_i`` weighing each sample ``(x_i, y_i)``.
  580. Weights are normalized to 1 if `density` is True. If `density` is
  581. False, the values of the returned histogram are equal to the sum of
  582. the weights belonging to the samples falling into each bin.
  583. Returns
  584. -------
  585. H : ndarray, shape(nx, ny)
  586. The bi-dimensional histogram of samples `x` and `y`. Values in `x`
  587. are histogrammed along the first dimension and values in `y` are
  588. histogrammed along the second dimension.
  589. xedges : ndarray, shape(nx+1,)
  590. The bin edges along the first dimension.
  591. yedges : ndarray, shape(ny+1,)
  592. The bin edges along the second dimension.
  593. See Also
  594. --------
  595. histogram : 1D histogram
  596. histogramdd : Multidimensional histogram
  597. Notes
  598. -----
  599. When `density` is True, then the returned histogram is the sample
  600. density, defined such that the sum over bins of the product
  601. ``bin_value * bin_area`` is 1.
  602. Please note that the histogram does not follow the Cartesian convention
  603. where `x` values are on the abscissa and `y` values on the ordinate
  604. axis. Rather, `x` is histogrammed along the first dimension of the
  605. array (vertical), and `y` along the second dimension of the array
  606. (horizontal). This ensures compatibility with `histogramdd`.
  607. Examples
  608. --------
  609. >>> import numpy as np
  610. >>> from matplotlib.image import NonUniformImage
  611. >>> import matplotlib.pyplot as plt
  612. Construct a 2-D histogram with variable bin width. First define the bin
  613. edges:
  614. >>> xedges = [0, 1, 3, 5]
  615. >>> yedges = [0, 2, 3, 4, 6]
  616. Next we create a histogram H with random bin content:
  617. >>> x = np.random.normal(2, 1, 100)
  618. >>> y = np.random.normal(1, 1, 100)
  619. >>> H, xedges, yedges = np.histogram2d(x, y, bins=(xedges, yedges))
  620. >>> # Histogram does not follow Cartesian convention (see Notes),
  621. >>> # therefore transpose H for visualization purposes.
  622. >>> H = H.T
  623. :func:`imshow <matplotlib.pyplot.imshow>` can only display square bins:
  624. >>> fig = plt.figure(figsize=(7, 3))
  625. >>> ax = fig.add_subplot(131, title='imshow: square bins')
  626. >>> plt.imshow(H, interpolation='nearest', origin='lower',
  627. ... extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]])
  628. <matplotlib.image.AxesImage object at 0x...>
  629. :func:`pcolormesh <matplotlib.pyplot.pcolormesh>` can display actual edges:
  630. >>> ax = fig.add_subplot(132, title='pcolormesh: actual edges',
  631. ... aspect='equal')
  632. >>> X, Y = np.meshgrid(xedges, yedges)
  633. >>> ax.pcolormesh(X, Y, H)
  634. <matplotlib.collections.QuadMesh object at 0x...>
  635. :class:`NonUniformImage <matplotlib.image.NonUniformImage>` can be used to
  636. display actual bin edges with interpolation:
  637. >>> ax = fig.add_subplot(133, title='NonUniformImage: interpolated',
  638. ... aspect='equal', xlim=xedges[[0, -1]], ylim=yedges[[0, -1]])
  639. >>> im = NonUniformImage(ax, interpolation='bilinear')
  640. >>> xcenters = (xedges[:-1] + xedges[1:]) / 2
  641. >>> ycenters = (yedges[:-1] + yedges[1:]) / 2
  642. >>> im.set_data(xcenters, ycenters, H)
  643. >>> ax.add_image(im)
  644. >>> plt.show()
  645. It is also possible to construct a 2-D histogram without specifying bin
  646. edges:
  647. >>> # Generate non-symmetric test data
  648. >>> n = 10000
  649. >>> x = np.linspace(1, 100, n)
  650. >>> y = 2*np.log(x) + np.random.rand(n) - 0.5
  651. >>> # Compute 2d histogram. Note the order of x/y and xedges/yedges
  652. >>> H, yedges, xedges = np.histogram2d(y, x, bins=20)
  653. Now we can plot the histogram using
  654. :func:`pcolormesh <matplotlib.pyplot.pcolormesh>`, and a
  655. :func:`hexbin <matplotlib.pyplot.hexbin>` for comparison.
  656. >>> # Plot histogram using pcolormesh
  657. >>> fig, (ax1, ax2) = plt.subplots(ncols=2, sharey=True)
  658. >>> ax1.pcolormesh(xedges, yedges, H, cmap='rainbow')
  659. >>> ax1.plot(x, 2*np.log(x), 'k-')
  660. >>> ax1.set_xlim(x.min(), x.max())
  661. >>> ax1.set_ylim(y.min(), y.max())
  662. >>> ax1.set_xlabel('x')
  663. >>> ax1.set_ylabel('y')
  664. >>> ax1.set_title('histogram2d')
  665. >>> ax1.grid()
  666. >>> # Create hexbin plot for comparison
  667. >>> ax2.hexbin(x, y, gridsize=20, cmap='rainbow')
  668. >>> ax2.plot(x, 2*np.log(x), 'k-')
  669. >>> ax2.set_title('hexbin')
  670. >>> ax2.set_xlim(x.min(), x.max())
  671. >>> ax2.set_xlabel('x')
  672. >>> ax2.grid()
  673. >>> plt.show()
  674. """
  675. from numpy import histogramdd
  676. if len(x) != len(y):
  677. raise ValueError('x and y must have the same length.')
  678. try:
  679. N = len(bins)
  680. except TypeError:
  681. N = 1
  682. if N not in {1, 2}:
  683. xedges = yedges = asarray(bins)
  684. bins = [xedges, yedges]
  685. hist, edges = histogramdd([x, y], bins, range, density, weights)
  686. return hist, edges[0], edges[1]
  687. @set_module('numpy')
  688. def mask_indices(n, mask_func, k=0):
  689. """
  690. Return the indices to access (n, n) arrays, given a masking function.
  691. Assume `mask_func` is a function that, for a square array a of size
  692. ``(n, n)`` with a possible offset argument `k`, when called as
  693. ``mask_func(a, k)`` returns a new array with zeros in certain locations
  694. (functions like `triu` or `tril` do precisely this). Then this function
  695. returns the indices where the non-zero values would be located.
  696. Parameters
  697. ----------
  698. n : int
  699. The returned indices will be valid to access arrays of shape (n, n).
  700. mask_func : callable
  701. A function whose call signature is similar to that of `triu`, `tril`.
  702. That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`.
  703. `k` is an optional argument to the function.
  704. k : scalar
  705. An optional argument which is passed through to `mask_func`. Functions
  706. like `triu`, `tril` take a second argument that is interpreted as an
  707. offset.
  708. Returns
  709. -------
  710. indices : tuple of arrays.
  711. The `n` arrays of indices corresponding to the locations where
  712. ``mask_func(np.ones((n, n)), k)`` is True.
  713. See Also
  714. --------
  715. triu, tril, triu_indices, tril_indices
  716. Examples
  717. --------
  718. >>> import numpy as np
  719. These are the indices that would allow you to access the upper triangular
  720. part of any 3x3 array:
  721. >>> iu = np.mask_indices(3, np.triu)
  722. For example, if `a` is a 3x3 array:
  723. >>> a = np.arange(9).reshape(3, 3)
  724. >>> a
  725. array([[0, 1, 2],
  726. [3, 4, 5],
  727. [6, 7, 8]])
  728. >>> a[iu]
  729. array([0, 1, 2, 4, 5, 8])
  730. An offset can be passed also to the masking function. This gets us the
  731. indices starting on the first diagonal right of the main one:
  732. >>> iu1 = np.mask_indices(3, np.triu, 1)
  733. with which we now extract only three elements:
  734. >>> a[iu1]
  735. array([1, 2, 5])
  736. """
  737. m = ones((n, n), int)
  738. a = mask_func(m, k)
  739. return nonzero(a != 0)
  740. @set_module('numpy')
  741. def tril_indices(n, k=0, m=None):
  742. """
  743. Return the indices for the lower-triangle of an (n, m) array.
  744. Parameters
  745. ----------
  746. n : int
  747. The row dimension of the arrays for which the returned
  748. indices will be valid.
  749. k : int, optional
  750. Diagonal offset (see `tril` for details).
  751. m : int, optional
  752. The column dimension of the arrays for which the returned
  753. arrays will be valid.
  754. By default `m` is taken equal to `n`.
  755. Returns
  756. -------
  757. inds : tuple of arrays
  758. The row and column indices, respectively. The row indices are sorted
  759. in non-decreasing order, and the corresponding column indices are
  760. strictly increasing for each row.
  761. See also
  762. --------
  763. triu_indices : similar function, for upper-triangular.
  764. mask_indices : generic function accepting an arbitrary mask function.
  765. tril, triu
  766. Examples
  767. --------
  768. >>> import numpy as np
  769. Compute two different sets of indices to access 4x4 arrays, one for the
  770. lower triangular part starting at the main diagonal, and one starting two
  771. diagonals further right:
  772. >>> il1 = np.tril_indices(4)
  773. >>> il1
  774. (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3]))
  775. Note that row indices (first array) are non-decreasing, and the corresponding
  776. column indices (second array) are strictly increasing for each row.
  777. Here is how they can be used with a sample array:
  778. >>> a = np.arange(16).reshape(4, 4)
  779. >>> a
  780. array([[ 0, 1, 2, 3],
  781. [ 4, 5, 6, 7],
  782. [ 8, 9, 10, 11],
  783. [12, 13, 14, 15]])
  784. Both for indexing:
  785. >>> a[il1]
  786. array([ 0, 4, 5, ..., 13, 14, 15])
  787. And for assigning values:
  788. >>> a[il1] = -1
  789. >>> a
  790. array([[-1, 1, 2, 3],
  791. [-1, -1, 6, 7],
  792. [-1, -1, -1, 11],
  793. [-1, -1, -1, -1]])
  794. These cover almost the whole array (two diagonals right of the main one):
  795. >>> il2 = np.tril_indices(4, 2)
  796. >>> a[il2] = -10
  797. >>> a
  798. array([[-10, -10, -10, 3],
  799. [-10, -10, -10, -10],
  800. [-10, -10, -10, -10],
  801. [-10, -10, -10, -10]])
  802. """
  803. tri_ = tri(n, m, k=k, dtype=bool)
  804. return tuple(broadcast_to(inds, tri_.shape)[tri_]
  805. for inds in indices(tri_.shape, sparse=True))
  806. def _trilu_indices_form_dispatcher(arr, k=None):
  807. return (arr,)
  808. @array_function_dispatch(_trilu_indices_form_dispatcher)
  809. def tril_indices_from(arr, k=0):
  810. """
  811. Return the indices for the lower-triangle of arr.
  812. See `tril_indices` for full details.
  813. Parameters
  814. ----------
  815. arr : array_like
  816. The indices will be valid for square arrays whose dimensions are
  817. the same as arr.
  818. k : int, optional
  819. Diagonal offset (see `tril` for details).
  820. Examples
  821. --------
  822. >>> import numpy as np
  823. Create a 4 by 4 array
  824. >>> a = np.arange(16).reshape(4, 4)
  825. >>> a
  826. array([[ 0, 1, 2, 3],
  827. [ 4, 5, 6, 7],
  828. [ 8, 9, 10, 11],
  829. [12, 13, 14, 15]])
  830. Pass the array to get the indices of the lower triangular elements.
  831. >>> trili = np.tril_indices_from(a)
  832. >>> trili
  833. (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3]))
  834. >>> a[trili]
  835. array([ 0, 4, 5, 8, 9, 10, 12, 13, 14, 15])
  836. This is syntactic sugar for tril_indices().
  837. >>> np.tril_indices(a.shape[0])
  838. (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3]))
  839. Use the `k` parameter to return the indices for the lower triangular array
  840. up to the k-th diagonal.
  841. >>> trili1 = np.tril_indices_from(a, k=1)
  842. >>> a[trili1]
  843. array([ 0, 1, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15])
  844. See Also
  845. --------
  846. tril_indices, tril, triu_indices_from
  847. """
  848. if arr.ndim != 2:
  849. raise ValueError("input array must be 2-d")
  850. return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1])
  851. @set_module('numpy')
  852. def triu_indices(n, k=0, m=None):
  853. """
  854. Return the indices for the upper-triangle of an (n, m) array.
  855. Parameters
  856. ----------
  857. n : int
  858. The size of the arrays for which the returned indices will
  859. be valid.
  860. k : int, optional
  861. Diagonal offset (see `triu` for details).
  862. m : int, optional
  863. The column dimension of the arrays for which the returned
  864. arrays will be valid.
  865. By default `m` is taken equal to `n`.
  866. Returns
  867. -------
  868. inds : tuple, shape(2) of ndarrays, shape(`n`)
  869. The row and column indices, respectively. The row indices are sorted
  870. in non-decreasing order, and the corresponding column indices are
  871. strictly increasing for each row.
  872. See also
  873. --------
  874. tril_indices : similar function, for lower-triangular.
  875. mask_indices : generic function accepting an arbitrary mask function.
  876. triu, tril
  877. Examples
  878. --------
  879. >>> import numpy as np
  880. Compute two different sets of indices to access 4x4 arrays, one for the
  881. upper triangular part starting at the main diagonal, and one starting two
  882. diagonals further right:
  883. >>> iu1 = np.triu_indices(4)
  884. >>> iu1
  885. (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3]))
  886. Note that row indices (first array) are non-decreasing, and the corresponding
  887. column indices (second array) are strictly increasing for each row.
  888. Here is how they can be used with a sample array:
  889. >>> a = np.arange(16).reshape(4, 4)
  890. >>> a
  891. array([[ 0, 1, 2, 3],
  892. [ 4, 5, 6, 7],
  893. [ 8, 9, 10, 11],
  894. [12, 13, 14, 15]])
  895. Both for indexing:
  896. >>> a[iu1]
  897. array([ 0, 1, 2, ..., 10, 11, 15])
  898. And for assigning values:
  899. >>> a[iu1] = -1
  900. >>> a
  901. array([[-1, -1, -1, -1],
  902. [ 4, -1, -1, -1],
  903. [ 8, 9, -1, -1],
  904. [12, 13, 14, -1]])
  905. These cover only a small part of the whole array (two diagonals right
  906. of the main one):
  907. >>> iu2 = np.triu_indices(4, 2)
  908. >>> a[iu2] = -10
  909. >>> a
  910. array([[ -1, -1, -10, -10],
  911. [ 4, -1, -1, -10],
  912. [ 8, 9, -1, -1],
  913. [ 12, 13, 14, -1]])
  914. """
  915. tri_ = ~tri(n, m, k=k - 1, dtype=bool)
  916. return tuple(broadcast_to(inds, tri_.shape)[tri_]
  917. for inds in indices(tri_.shape, sparse=True))
  918. @array_function_dispatch(_trilu_indices_form_dispatcher)
  919. def triu_indices_from(arr, k=0):
  920. """
  921. Return the indices for the upper-triangle of arr.
  922. See `triu_indices` for full details.
  923. Parameters
  924. ----------
  925. arr : ndarray, shape(N, N)
  926. The indices will be valid for square arrays.
  927. k : int, optional
  928. Diagonal offset (see `triu` for details).
  929. Returns
  930. -------
  931. triu_indices_from : tuple, shape(2) of ndarray, shape(N)
  932. Indices for the upper-triangle of `arr`.
  933. Examples
  934. --------
  935. >>> import numpy as np
  936. Create a 4 by 4 array
  937. >>> a = np.arange(16).reshape(4, 4)
  938. >>> a
  939. array([[ 0, 1, 2, 3],
  940. [ 4, 5, 6, 7],
  941. [ 8, 9, 10, 11],
  942. [12, 13, 14, 15]])
  943. Pass the array to get the indices of the upper triangular elements.
  944. >>> triui = np.triu_indices_from(a)
  945. >>> triui
  946. (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3]))
  947. >>> a[triui]
  948. array([ 0, 1, 2, 3, 5, 6, 7, 10, 11, 15])
  949. This is syntactic sugar for triu_indices().
  950. >>> np.triu_indices(a.shape[0])
  951. (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3]))
  952. Use the `k` parameter to return the indices for the upper triangular array
  953. from the k-th diagonal.
  954. >>> triuim1 = np.triu_indices_from(a, k=1)
  955. >>> a[triuim1]
  956. array([ 1, 2, 3, 6, 7, 11])
  957. See Also
  958. --------
  959. triu_indices, triu, tril_indices_from
  960. """
  961. if arr.ndim != 2:
  962. raise ValueError("input array must be 2-d")
  963. return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1])