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