Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1""" Basic functions for manipulating 2d arrays
3"""
4import functools
5import operator
7from numpy._core._multiarray_umath import _array_converter
8from numpy._core.numeric import (
9 asanyarray, arange, zeros, greater_equal, multiply, ones,
10 asarray, where, int8, int16, int32, int64, intp, empty, promote_types,
11 diagonal, nonzero, indices
12 )
13from numpy._core.overrides import finalize_array_function_like, set_module
14from numpy._core import overrides
15from numpy._core import iinfo
16from numpy.lib._stride_tricks_impl import broadcast_to
19__all__ = [
20 'diag', 'diagflat', 'eye', 'fliplr', 'flipud', 'tri', 'triu',
21 'tril', 'vander', 'histogram2d', 'mask_indices', 'tril_indices',
22 'tril_indices_from', 'triu_indices', 'triu_indices_from', ]
25array_function_dispatch = functools.partial(
26 overrides.array_function_dispatch, module='numpy')
29i1 = iinfo(int8)
30i2 = iinfo(int16)
31i4 = iinfo(int32)
34def _min_int(low, high):
35 """ get small int that fits the range """
36 if high <= i1.max and low >= i1.min:
37 return int8
38 if high <= i2.max and low >= i2.min:
39 return int16
40 if high <= i4.max and low >= i4.min:
41 return int32
42 return int64
45def _flip_dispatcher(m):
46 return (m,)
49@array_function_dispatch(_flip_dispatcher)
50def fliplr(m):
51 """
52 Reverse the order of elements along axis 1 (left/right).
54 For a 2-D array, this flips the entries in each row in the left/right
55 direction. Columns are preserved, but appear in a different order than
56 before.
58 Parameters
59 ----------
60 m : array_like
61 Input array, must be at least 2-D.
63 Returns
64 -------
65 f : ndarray
66 A view of `m` with the columns reversed. Since a view
67 is returned, this operation is :math:`\\mathcal O(1)`.
69 See Also
70 --------
71 flipud : Flip array in the up/down direction.
72 flip : Flip array in one or more dimensions.
73 rot90 : Rotate array counterclockwise.
75 Notes
76 -----
77 Equivalent to ``m[:,::-1]`` or ``np.flip(m, axis=1)``.
78 Requires the array to be at least 2-D.
80 Examples
81 --------
82 >>> import numpy as np
83 >>> A = np.diag([1.,2.,3.])
84 >>> A
85 array([[1., 0., 0.],
86 [0., 2., 0.],
87 [0., 0., 3.]])
88 >>> np.fliplr(A)
89 array([[0., 0., 1.],
90 [0., 2., 0.],
91 [3., 0., 0.]])
93 >>> rng = np.random.default_rng()
94 >>> A = rng.normal(size=(2,3,5))
95 >>> np.all(np.fliplr(A) == A[:,::-1,...])
96 True
98 """
99 m = asanyarray(m)
100 if m.ndim < 2:
101 raise ValueError("Input must be >= 2-d.")
102 return m[:, ::-1]
105@array_function_dispatch(_flip_dispatcher)
106def flipud(m):
107 """
108 Reverse the order of elements along axis 0 (up/down).
110 For a 2-D array, this flips the entries in each column in the up/down
111 direction. Rows are preserved, but appear in a different order than before.
113 Parameters
114 ----------
115 m : array_like
116 Input array.
118 Returns
119 -------
120 out : array_like
121 A view of `m` with the rows reversed. Since a view is
122 returned, this operation is :math:`\\mathcal O(1)`.
124 See Also
125 --------
126 fliplr : Flip array in the left/right direction.
127 flip : Flip array in one or more dimensions.
128 rot90 : Rotate array counterclockwise.
130 Notes
131 -----
132 Equivalent to ``m[::-1, ...]`` or ``np.flip(m, axis=0)``.
133 Requires the array to be at least 1-D.
135 Examples
136 --------
137 >>> import numpy as np
138 >>> A = np.diag([1.0, 2, 3])
139 >>> A
140 array([[1., 0., 0.],
141 [0., 2., 0.],
142 [0., 0., 3.]])
143 >>> np.flipud(A)
144 array([[0., 0., 3.],
145 [0., 2., 0.],
146 [1., 0., 0.]])
148 >>> rng = np.random.default_rng()
149 >>> A = rng.normal(size=(2,3,5))
150 >>> np.all(np.flipud(A) == A[::-1,...])
151 True
153 >>> np.flipud([1,2])
154 array([2, 1])
156 """
157 m = asanyarray(m)
158 if m.ndim < 1:
159 raise ValueError("Input must be >= 1-d.")
160 return m[::-1, ...]
163@finalize_array_function_like
164@set_module('numpy')
165def eye(N, M=None, k=0, dtype=float, order='C', *, device=None, like=None):
166 """
167 Return a 2-D array with ones on the diagonal and zeros elsewhere.
169 Parameters
170 ----------
171 N : int
172 Number of rows in the output.
173 M : int, optional
174 Number of columns in the output. If None, defaults to `N`.
175 k : int, optional
176 Index of the diagonal: 0 (the default) refers to the main diagonal,
177 a positive value refers to an upper diagonal, and a negative value
178 to a lower diagonal.
179 dtype : data-type, optional
180 Data-type of the returned array.
181 order : {'C', 'F'}, optional
182 Whether the output should be stored in row-major (C-style) or
183 column-major (Fortran-style) order in memory.
184 device : str, optional
185 The device on which to place the created array. Default: None.
186 For Array-API interoperability only, so must be ``"cpu"`` if passed.
188 .. versionadded:: 2.0.0
189 ${ARRAY_FUNCTION_LIKE}
191 .. versionadded:: 1.20.0
193 Returns
194 -------
195 I : ndarray of shape (N,M)
196 An array where all elements are equal to zero, except for the `k`-th
197 diagonal, whose values are equal to one.
199 See Also
200 --------
201 identity : (almost) equivalent function
202 diag : diagonal 2-D array from a 1-D array specified by the user.
204 Examples
205 --------
206 >>> import numpy as np
207 >>> np.eye(2, dtype=int)
208 array([[1, 0],
209 [0, 1]])
210 >>> np.eye(3, k=1)
211 array([[0., 1., 0.],
212 [0., 0., 1.],
213 [0., 0., 0.]])
215 """
216 if like is not None:
217 return _eye_with_like(
218 like, N, M=M, k=k, dtype=dtype, order=order, device=device
219 )
220 if M is None:
221 M = N
222 m = zeros((N, M), dtype=dtype, order=order, device=device)
223 if k >= M:
224 return m
225 # Ensure M and k are integers, so we don't get any surprise casting
226 # results in the expressions `M-k` and `M+1` used below. This avoids
227 # a problem with inputs with type (for example) np.uint64.
228 M = operator.index(M)
229 k = operator.index(k)
230 if k >= 0:
231 i = k
232 else:
233 i = (-k) * M
234 m[:M-k].flat[i::M+1] = 1
235 return m
238_eye_with_like = array_function_dispatch()(eye)
241def _diag_dispatcher(v, k=None):
242 return (v,)
245@array_function_dispatch(_diag_dispatcher)
246def diag(v, k=0):
247 """
248 Extract a diagonal or construct a diagonal array.
250 See the more detailed documentation for ``numpy.diagonal`` if you use this
251 function to extract a diagonal and wish to write to the resulting array;
252 whether it returns a copy or a view depends on what version of numpy you
253 are using.
255 Parameters
256 ----------
257 v : array_like
258 If `v` is a 2-D array, return a copy of its `k`-th diagonal.
259 If `v` is a 1-D array, return a 2-D array with `v` on the `k`-th
260 diagonal.
261 k : int, optional
262 Diagonal in question. The default is 0. Use `k>0` for diagonals
263 above the main diagonal, and `k<0` for diagonals below the main
264 diagonal.
266 Returns
267 -------
268 out : ndarray
269 The extracted diagonal or constructed diagonal array.
271 See Also
272 --------
273 diagonal : Return specified diagonals.
274 diagflat : Create a 2-D array with the flattened input as a diagonal.
275 trace : Sum along diagonals.
276 triu : Upper triangle of an array.
277 tril : Lower triangle of an array.
279 Examples
280 --------
281 >>> import numpy as np
282 >>> x = np.arange(9).reshape((3,3))
283 >>> x
284 array([[0, 1, 2],
285 [3, 4, 5],
286 [6, 7, 8]])
288 >>> np.diag(x)
289 array([0, 4, 8])
290 >>> np.diag(x, k=1)
291 array([1, 5])
292 >>> np.diag(x, k=-1)
293 array([3, 7])
295 >>> np.diag(np.diag(x))
296 array([[0, 0, 0],
297 [0, 4, 0],
298 [0, 0, 8]])
300 """
301 v = asanyarray(v)
302 s = v.shape
303 if len(s) == 1:
304 n = s[0]+abs(k)
305 res = zeros((n, n), v.dtype)
306 if k >= 0:
307 i = k
308 else:
309 i = (-k) * n
310 res[:n-k].flat[i::n+1] = v
311 return res
312 elif len(s) == 2:
313 return diagonal(v, k)
314 else:
315 raise ValueError("Input must be 1- or 2-d.")
318@array_function_dispatch(_diag_dispatcher)
319def diagflat(v, k=0):
320 """
321 Create a two-dimensional array with the flattened input as a diagonal.
323 Parameters
324 ----------
325 v : array_like
326 Input data, which is flattened and set as the `k`-th
327 diagonal of the output.
328 k : int, optional
329 Diagonal to set; 0, the default, corresponds to the "main" diagonal,
330 a positive (negative) `k` giving the number of the diagonal above
331 (below) the main.
333 Returns
334 -------
335 out : ndarray
336 The 2-D output array.
338 See Also
339 --------
340 diag : MATLAB work-alike for 1-D and 2-D arrays.
341 diagonal : Return specified diagonals.
342 trace : Sum along diagonals.
344 Examples
345 --------
346 >>> import numpy as np
347 >>> np.diagflat([[1,2], [3,4]])
348 array([[1, 0, 0, 0],
349 [0, 2, 0, 0],
350 [0, 0, 3, 0],
351 [0, 0, 0, 4]])
353 >>> np.diagflat([1,2], 1)
354 array([[0, 1, 0],
355 [0, 0, 2],
356 [0, 0, 0]])
358 """
359 conv = _array_converter(v)
360 v, = conv.as_arrays(subok=False)
361 v = v.ravel()
362 s = len(v)
363 n = s + abs(k)
364 res = zeros((n, n), v.dtype)
365 if (k >= 0):
366 i = arange(0, n-k, dtype=intp)
367 fi = i+k+i*n
368 else:
369 i = arange(0, n+k, dtype=intp)
370 fi = i+(i-k)*n
371 res.flat[fi] = v
373 return conv.wrap(res)
376@finalize_array_function_like
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.
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}
397 .. versionadded:: 1.20.0
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.
405 Examples
406 --------
407 >>> import numpy as np
408 >>> np.tri(3, 5, 2, dtype=int)
409 array([[1, 1, 1, 0, 0],
410 [1, 1, 1, 1, 0],
411 [1, 1, 1, 1, 1]])
413 >>> np.tri(3, 5, -1)
414 array([[0., 0., 0., 0., 0.],
415 [1., 0., 0., 0., 0.],
416 [1., 1., 0., 0., 0.]])
418 """
419 if like is not None:
420 return _tri_with_like(like, N, M=M, k=k, dtype=dtype)
422 if M is None:
423 M = N
425 m = greater_equal.outer(arange(N, dtype=_min_int(0, N)),
426 arange(-k, M-k, dtype=_min_int(-k, M - k)))
428 # Avoid making a copy if the requested type is already bool
429 m = m.astype(dtype, copy=False)
431 return m
434_tri_with_like = array_function_dispatch()(tri)
437def _trilu_dispatcher(m, k=None):
438 return (m,)
441@array_function_dispatch(_trilu_dispatcher)
442def tril(m, k=0):
443 """
444 Lower triangle of an array.
446 Return a copy of an array with elements above the `k`-th diagonal zeroed.
447 For arrays with ``ndim`` exceeding 2, `tril` will apply to the final two
448 axes.
450 Parameters
451 ----------
452 m : array_like, shape (..., M, N)
453 Input array.
454 k : int, optional
455 Diagonal above which to zero elements. `k = 0` (the default) is the
456 main diagonal, `k < 0` is below it and `k > 0` is above.
458 Returns
459 -------
460 tril : ndarray, shape (..., M, N)
461 Lower triangle of `m`, of same shape and data-type as `m`.
463 See Also
464 --------
465 triu : same thing, only for the upper triangle
467 Examples
468 --------
469 >>> import numpy as np
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]])
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]]])
490 """
491 m = asanyarray(m)
492 mask = tri(*m.shape[-2:], k=k, dtype=bool)
494 return where(mask, m, zeros(1, m.dtype))
497@array_function_dispatch(_trilu_dispatcher)
498def triu(m, k=0):
499 """
500 Upper triangle of an array.
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.
506 Please refer to the documentation for `tril` for further details.
508 See Also
509 --------
510 tril : lower triangle of an array
512 Examples
513 --------
514 >>> import numpy as np
515 >>> np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1)
516 array([[ 1, 2, 3],
517 [ 4, 5, 6],
518 [ 0, 8, 9],
519 [ 0, 0, 12]])
521 >>> np.triu(np.arange(3*4*5).reshape(3, 4, 5))
522 array([[[ 0, 1, 2, 3, 4],
523 [ 0, 6, 7, 8, 9],
524 [ 0, 0, 12, 13, 14],
525 [ 0, 0, 0, 18, 19]],
526 [[20, 21, 22, 23, 24],
527 [ 0, 26, 27, 28, 29],
528 [ 0, 0, 32, 33, 34],
529 [ 0, 0, 0, 38, 39]],
530 [[40, 41, 42, 43, 44],
531 [ 0, 46, 47, 48, 49],
532 [ 0, 0, 52, 53, 54],
533 [ 0, 0, 0, 58, 59]]])
535 """
536 m = asanyarray(m)
537 mask = tri(*m.shape[-2:], k=k-1, dtype=bool)
539 return where(mask, zeros(1, m.dtype), m)
542def _vander_dispatcher(x, N=None, increasing=None):
543 return (x,)
546# Originally borrowed from John Hunter and matplotlib
547@array_function_dispatch(_vander_dispatcher)
548def vander(x, N=None, increasing=False):
549 """
550 Generate a Vandermonde matrix.
552 The columns of the output matrix are powers of the input vector. The
553 order of the powers is determined by the `increasing` boolean argument.
554 Specifically, when `increasing` is False, the `i`-th output column is
555 the input vector raised element-wise to the power of ``N - i - 1``. Such
556 a matrix with a geometric progression in each row is named for Alexandre-
557 Theophile Vandermonde.
559 Parameters
560 ----------
561 x : array_like
562 1-D input array.
563 N : int, optional
564 Number of columns in the output. If `N` is not specified, a square
565 array is returned (``N = len(x)``).
566 increasing : bool, optional
567 Order of the powers of the columns. If True, the powers increase
568 from left to right, if False (the default) they are reversed.
570 Returns
571 -------
572 out : ndarray
573 Vandermonde matrix. If `increasing` is False, the first column is
574 ``x^(N-1)``, the second ``x^(N-2)`` and so forth. If `increasing` is
575 True, the columns are ``x^0, x^1, ..., x^(N-1)``.
577 See Also
578 --------
579 polynomial.polynomial.polyvander
581 Examples
582 --------
583 >>> import numpy as np
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]])
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]])
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]])
610 The determinant of a square Vandermonde matrix is the product
611 of the differences between the values of the input vector:
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
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)
625 v = empty((len(x), N), dtype=promote_types(x.dtype, int))
626 tmp = v[:, ::-1] if not increasing else v
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)
634 return v
637def _histogram2d_dispatcher(x, y, bins=None, range=None, density=None,
638 weights=None):
639 yield x
640 yield y
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
652 yield weights
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.
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:
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.
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.
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.
707 See Also
708 --------
709 histogram : 1D histogram
710 histogramdd : Multidimensional histogram
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.
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`.
724 Examples
725 --------
726 >>> import numpy as np
727 >>> from matplotlib.image import NonUniformImage
728 >>> import matplotlib.pyplot as plt
730 Construct a 2-D histogram with variable bin width. First define the bin
731 edges:
733 >>> xedges = [0, 1, 3, 5]
734 >>> yedges = [0, 2, 3, 4, 6]
736 Next we create a histogram H with random bin content:
738 >>> x = np.random.normal(2, 1, 100)
739 >>> y = np.random.normal(1, 1, 100)
740 >>> H, xedges, yedges = np.histogram2d(x, y, bins=(xedges, yedges))
741 >>> # Histogram does not follow Cartesian convention (see Notes),
742 >>> # therefore transpose H for visualization purposes.
743 >>> H = H.T
745 :func:`imshow <matplotlib.pyplot.imshow>` can only display square bins:
747 >>> fig = plt.figure(figsize=(7, 3))
748 >>> ax = fig.add_subplot(131, title='imshow: square bins')
749 >>> plt.imshow(H, interpolation='nearest', origin='lower',
750 ... extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]])
751 <matplotlib.image.AxesImage object at 0x...>
753 :func:`pcolormesh <matplotlib.pyplot.pcolormesh>` can display actual edges:
755 >>> ax = fig.add_subplot(132, title='pcolormesh: actual edges',
756 ... aspect='equal')
757 >>> X, Y = np.meshgrid(xedges, yedges)
758 >>> ax.pcolormesh(X, Y, H)
759 <matplotlib.collections.QuadMesh object at 0x...>
761 :class:`NonUniformImage <matplotlib.image.NonUniformImage>` can be used to
762 display actual bin edges with interpolation:
764 >>> ax = fig.add_subplot(133, title='NonUniformImage: interpolated',
765 ... aspect='equal', xlim=xedges[[0, -1]], ylim=yedges[[0, -1]])
766 >>> im = NonUniformImage(ax, interpolation='bilinear')
767 >>> xcenters = (xedges[:-1] + xedges[1:]) / 2
768 >>> ycenters = (yedges[:-1] + yedges[1:]) / 2
769 >>> im.set_data(xcenters, ycenters, H)
770 >>> ax.add_image(im)
771 >>> plt.show()
773 It is also possible to construct a 2-D histogram without specifying bin
774 edges:
776 >>> # Generate non-symmetric test data
777 >>> n = 10000
778 >>> x = np.linspace(1, 100, n)
779 >>> y = 2*np.log(x) + np.random.rand(n) - 0.5
780 >>> # Compute 2d histogram. Note the order of x/y and xedges/yedges
781 >>> H, yedges, xedges = np.histogram2d(y, x, bins=20)
783 Now we can plot the histogram using
784 :func:`pcolormesh <matplotlib.pyplot.pcolormesh>`, and a
785 :func:`hexbin <matplotlib.pyplot.hexbin>` for comparison.
787 >>> # Plot histogram using pcolormesh
788 >>> fig, (ax1, ax2) = plt.subplots(ncols=2, sharey=True)
789 >>> ax1.pcolormesh(xedges, yedges, H, cmap='rainbow')
790 >>> ax1.plot(x, 2*np.log(x), 'k-')
791 >>> ax1.set_xlim(x.min(), x.max())
792 >>> ax1.set_ylim(y.min(), y.max())
793 >>> ax1.set_xlabel('x')
794 >>> ax1.set_ylabel('y')
795 >>> ax1.set_title('histogram2d')
796 >>> ax1.grid()
798 >>> # Create hexbin plot for comparison
799 >>> ax2.hexbin(x, y, gridsize=20, cmap='rainbow')
800 >>> ax2.plot(x, 2*np.log(x), 'k-')
801 >>> ax2.set_title('hexbin')
802 >>> ax2.set_xlim(x.min(), x.max())
803 >>> ax2.set_xlabel('x')
804 >>> ax2.grid()
806 >>> plt.show()
807 """
808 from numpy import histogramdd
810 if len(x) != len(y):
811 raise ValueError('x and y must have the same length.')
813 try:
814 N = len(bins)
815 except TypeError:
816 N = 1
818 if N != 1 and N != 2:
819 xedges = yedges = asarray(bins)
820 bins = [xedges, yedges]
821 hist, edges = histogramdd([x, y], bins, range, density, weights)
822 return hist, edges[0], edges[1]
825@set_module('numpy')
826def mask_indices(n, mask_func, k=0):
827 """
828 Return the indices to access (n, n) arrays, given a masking function.
830 Assume `mask_func` is a function that, for a square array a of size
831 ``(n, n)`` with a possible offset argument `k`, when called as
832 ``mask_func(a, k)`` returns a new array with zeros in certain locations
833 (functions like `triu` or `tril` do precisely this). Then this function
834 returns the indices where the non-zero values would be located.
836 Parameters
837 ----------
838 n : int
839 The returned indices will be valid to access arrays of shape (n, n).
840 mask_func : callable
841 A function whose call signature is similar to that of `triu`, `tril`.
842 That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`.
843 `k` is an optional argument to the function.
844 k : scalar
845 An optional argument which is passed through to `mask_func`. Functions
846 like `triu`, `tril` take a second argument that is interpreted as an
847 offset.
849 Returns
850 -------
851 indices : tuple of arrays.
852 The `n` arrays of indices corresponding to the locations where
853 ``mask_func(np.ones((n, n)), k)`` is True.
855 See Also
856 --------
857 triu, tril, triu_indices, tril_indices
859 Examples
860 --------
861 >>> import numpy as np
863 These are the indices that would allow you to access the upper triangular
864 part of any 3x3 array:
866 >>> iu = np.mask_indices(3, np.triu)
868 For example, if `a` is a 3x3 array:
870 >>> a = np.arange(9).reshape(3, 3)
871 >>> a
872 array([[0, 1, 2],
873 [3, 4, 5],
874 [6, 7, 8]])
875 >>> a[iu]
876 array([0, 1, 2, 4, 5, 8])
878 An offset can be passed also to the masking function. This gets us the
879 indices starting on the first diagonal right of the main one:
881 >>> iu1 = np.mask_indices(3, np.triu, 1)
883 with which we now extract only three elements:
885 >>> a[iu1]
886 array([1, 2, 5])
888 """
889 m = ones((n, n), int)
890 a = mask_func(m, k)
891 return nonzero(a != 0)
894@set_module('numpy')
895def tril_indices(n, k=0, m=None):
896 """
897 Return the indices for the lower-triangle of an (n, m) array.
899 Parameters
900 ----------
901 n : int
902 The row dimension of the arrays for which the returned
903 indices will be valid.
904 k : int, optional
905 Diagonal offset (see `tril` for details).
906 m : int, optional
907 The column dimension of the arrays for which the returned
908 arrays will be valid.
909 By default `m` is taken equal to `n`.
912 Returns
913 -------
914 inds : tuple of arrays
915 The row and column indices, respectively. The row indices are sorted
916 in non-decreasing order, and the correspdonding column indices are
917 strictly increasing for each row.
919 See also
920 --------
921 triu_indices : similar function, for upper-triangular.
922 mask_indices : generic function accepting an arbitrary mask function.
923 tril, triu
925 Examples
926 --------
927 >>> import numpy as np
929 Compute two different sets of indices to access 4x4 arrays, one for the
930 lower triangular part starting at the main diagonal, and one starting two
931 diagonals further right:
933 >>> il1 = np.tril_indices(4)
934 >>> il1
935 (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3]))
937 Note that row indices (first array) are non-decreasing, and the corresponding
938 column indices (second array) are strictly increasing for each row.
939 Here is how they can be used with a sample array:
941 >>> a = np.arange(16).reshape(4, 4)
942 >>> a
943 array([[ 0, 1, 2, 3],
944 [ 4, 5, 6, 7],
945 [ 8, 9, 10, 11],
946 [12, 13, 14, 15]])
948 Both for indexing:
950 >>> a[il1]
951 array([ 0, 4, 5, ..., 13, 14, 15])
953 And for assigning values:
955 >>> a[il1] = -1
956 >>> a
957 array([[-1, 1, 2, 3],
958 [-1, -1, 6, 7],
959 [-1, -1, -1, 11],
960 [-1, -1, -1, -1]])
962 These cover almost the whole array (two diagonals right of the main one):
964 >>> il2 = np.tril_indices(4, 2)
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]])
972 """
973 tri_ = tri(n, m, k=k, dtype=bool)
975 return tuple(broadcast_to(inds, tri_.shape)[tri_]
976 for inds in indices(tri_.shape, sparse=True))
979def _trilu_indices_form_dispatcher(arr, k=None):
980 return (arr,)
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.
988 See `tril_indices` for full details.
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).
998 Examples
999 --------
1000 >>> import numpy as np
1002 Create a 4 by 4 array
1004 >>> a = np.arange(16).reshape(4, 4)
1005 >>> a
1006 array([[ 0, 1, 2, 3],
1007 [ 4, 5, 6, 7],
1008 [ 8, 9, 10, 11],
1009 [12, 13, 14, 15]])
1011 Pass the array to get the indices of the lower triangular elements.
1013 >>> trili = np.tril_indices_from(a)
1014 >>> trili
1015 (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3]))
1017 >>> a[trili]
1018 array([ 0, 4, 5, 8, 9, 10, 12, 13, 14, 15])
1020 This is syntactic sugar for tril_indices().
1022 >>> np.tril_indices(a.shape[0])
1023 (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3]))
1025 Use the `k` parameter to return the indices for the lower triangular array
1026 up to the k-th diagonal.
1028 >>> trili1 = np.tril_indices_from(a, k=1)
1029 >>> a[trili1]
1030 array([ 0, 1, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15])
1032 See Also
1033 --------
1034 tril_indices, tril, triu_indices_from
1035 """
1036 if arr.ndim != 2:
1037 raise ValueError("input array must be 2-d")
1038 return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1])
1041@set_module('numpy')
1042def triu_indices(n, k=0, m=None):
1043 """
1044 Return the indices for the upper-triangle of an (n, m) array.
1046 Parameters
1047 ----------
1048 n : int
1049 The size of the arrays for which the returned indices will
1050 be valid.
1051 k : int, optional
1052 Diagonal offset (see `triu` for details).
1053 m : int, optional
1054 The column dimension of the arrays for which the returned
1055 arrays will be valid.
1056 By default `m` is taken equal to `n`.
1059 Returns
1060 -------
1061 inds : tuple, shape(2) of ndarrays, shape(`n`)
1062 The row and column indices, respectively. The row indices are sorted
1063 in non-decreasing order, and the correspdonding column indices are
1064 strictly increasing for each row.
1066 See also
1067 --------
1068 tril_indices : similar function, for lower-triangular.
1069 mask_indices : generic function accepting an arbitrary mask function.
1070 triu, tril
1072 Examples
1073 --------
1074 >>> import numpy as np
1076 Compute two different sets of indices to access 4x4 arrays, one for the
1077 upper triangular part starting at the main diagonal, and one starting two
1078 diagonals further right:
1080 >>> iu1 = np.triu_indices(4)
1081 >>> iu1
1082 (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3]))
1084 Note that row indices (first array) are non-decreasing, and the corresponding
1085 column indices (second array) are strictly increasing for each row.
1087 Here is how they can be used with a sample array:
1089 >>> a = np.arange(16).reshape(4, 4)
1090 >>> a
1091 array([[ 0, 1, 2, 3],
1092 [ 4, 5, 6, 7],
1093 [ 8, 9, 10, 11],
1094 [12, 13, 14, 15]])
1096 Both for indexing:
1098 >>> a[iu1]
1099 array([ 0, 1, 2, ..., 10, 11, 15])
1101 And for assigning values:
1103 >>> a[iu1] = -1
1104 >>> a
1105 array([[-1, -1, -1, -1],
1106 [ 4, -1, -1, -1],
1107 [ 8, 9, -1, -1],
1108 [12, 13, 14, -1]])
1110 These cover only a small part of the whole array (two diagonals right
1111 of the main one):
1113 >>> iu2 = np.triu_indices(4, 2)
1114 >>> a[iu2] = -10
1115 >>> a
1116 array([[ -1, -1, -10, -10],
1117 [ 4, -1, -1, -10],
1118 [ 8, 9, -1, -1],
1119 [ 12, 13, 14, -1]])
1121 """
1122 tri_ = ~tri(n, m, k=k - 1, dtype=bool)
1124 return tuple(broadcast_to(inds, tri_.shape)[tri_]
1125 for inds in indices(tri_.shape, sparse=True))
1128@array_function_dispatch(_trilu_indices_form_dispatcher)
1129def triu_indices_from(arr, k=0):
1130 """
1131 Return the indices for the upper-triangle of arr.
1133 See `triu_indices` for full details.
1135 Parameters
1136 ----------
1137 arr : ndarray, shape(N, N)
1138 The indices will be valid for square arrays.
1139 k : int, optional
1140 Diagonal offset (see `triu` for details).
1142 Returns
1143 -------
1144 triu_indices_from : tuple, shape(2) of ndarray, shape(N)
1145 Indices for the upper-triangle of `arr`.
1147 Examples
1148 --------
1149 >>> import numpy as np
1151 Create a 4 by 4 array
1153 >>> a = np.arange(16).reshape(4, 4)
1154 >>> a
1155 array([[ 0, 1, 2, 3],
1156 [ 4, 5, 6, 7],
1157 [ 8, 9, 10, 11],
1158 [12, 13, 14, 15]])
1160 Pass the array to get the indices of the upper triangular elements.
1162 >>> triui = np.triu_indices_from(a)
1163 >>> triui
1164 (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3]))
1166 >>> a[triui]
1167 array([ 0, 1, 2, 3, 5, 6, 7, 10, 11, 15])
1169 This is syntactic sugar for triu_indices().
1171 >>> np.triu_indices(a.shape[0])
1172 (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3]))
1174 Use the `k` parameter to return the indices for the upper triangular array
1175 from the k-th diagonal.
1177 >>> triuim1 = np.triu_indices_from(a, k=1)
1178 >>> a[triuim1]
1179 array([ 1, 2, 3, 6, 7, 11])
1182 See Also
1183 --------
1184 triu_indices, triu, tril_indices_from
1185 """
1186 if arr.ndim != 2:
1187 raise ValueError("input array must be 2-d")
1188 return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1])