twodim_base.py 27 KB

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