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