twodim_base.py 28 KB

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