Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/numpy/lib/twodim_base.py : 33%

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