Chris@87: """ Basic functions for manipulating 2d arrays Chris@87: Chris@87: """ Chris@87: from __future__ import division, absolute_import, print_function Chris@87: Chris@87: from numpy.core.numeric import ( Chris@87: asanyarray, arange, zeros, greater_equal, multiply, ones, asarray, Chris@87: where, int8, int16, int32, int64, empty, promote_types Chris@87: ) Chris@87: from numpy.core import iinfo Chris@87: Chris@87: Chris@87: __all__ = [ Chris@87: 'diag', 'diagflat', 'eye', 'fliplr', 'flipud', 'rot90', 'tri', 'triu', Chris@87: 'tril', 'vander', 'histogram2d', 'mask_indices', 'tril_indices', Chris@87: 'tril_indices_from', 'triu_indices', 'triu_indices_from', ] Chris@87: Chris@87: Chris@87: i1 = iinfo(int8) Chris@87: i2 = iinfo(int16) Chris@87: i4 = iinfo(int32) Chris@87: def _min_int(low, high): Chris@87: """ get small int that fits the range """ Chris@87: if high <= i1.max and low >= i1.min: Chris@87: return int8 Chris@87: if high <= i2.max and low >= i2.min: Chris@87: return int16 Chris@87: if high <= i4.max and low >= i4.min: Chris@87: return int32 Chris@87: return int64 Chris@87: Chris@87: Chris@87: def fliplr(m): Chris@87: """ Chris@87: Flip array in the left/right direction. Chris@87: Chris@87: Flip the entries in each row in the left/right direction. Chris@87: Columns are preserved, but appear in a different order than before. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: m : array_like Chris@87: Input array, must be at least 2-D. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: f : ndarray Chris@87: A view of `m` with the columns reversed. Since a view Chris@87: is returned, this operation is :math:`\\mathcal O(1)`. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: flipud : Flip array in the up/down direction. Chris@87: rot90 : Rotate array counterclockwise. Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: Equivalent to A[:,::-1]. Requires the array to be at least 2-D. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> A = np.diag([1.,2.,3.]) Chris@87: >>> A Chris@87: array([[ 1., 0., 0.], Chris@87: [ 0., 2., 0.], Chris@87: [ 0., 0., 3.]]) Chris@87: >>> np.fliplr(A) Chris@87: array([[ 0., 0., 1.], Chris@87: [ 0., 2., 0.], Chris@87: [ 3., 0., 0.]]) Chris@87: Chris@87: >>> A = np.random.randn(2,3,5) Chris@87: >>> np.all(np.fliplr(A)==A[:,::-1,...]) Chris@87: True Chris@87: Chris@87: """ Chris@87: m = asanyarray(m) Chris@87: if m.ndim < 2: Chris@87: raise ValueError("Input must be >= 2-d.") Chris@87: return m[:, ::-1] Chris@87: Chris@87: Chris@87: def flipud(m): Chris@87: """ Chris@87: Flip array in the up/down direction. Chris@87: Chris@87: Flip the entries in each column in the up/down direction. Chris@87: Rows are preserved, but appear in a different order than before. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: m : array_like Chris@87: Input array. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: out : array_like Chris@87: A view of `m` with the rows reversed. Since a view is Chris@87: returned, this operation is :math:`\\mathcal O(1)`. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: fliplr : Flip array in the left/right direction. Chris@87: rot90 : Rotate array counterclockwise. Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: Equivalent to ``A[::-1,...]``. Chris@87: Does not require the array to be two-dimensional. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> A = np.diag([1.0, 2, 3]) Chris@87: >>> A Chris@87: array([[ 1., 0., 0.], Chris@87: [ 0., 2., 0.], Chris@87: [ 0., 0., 3.]]) Chris@87: >>> np.flipud(A) Chris@87: array([[ 0., 0., 3.], Chris@87: [ 0., 2., 0.], Chris@87: [ 1., 0., 0.]]) Chris@87: Chris@87: >>> A = np.random.randn(2,3,5) Chris@87: >>> np.all(np.flipud(A)==A[::-1,...]) Chris@87: True Chris@87: Chris@87: >>> np.flipud([1,2]) Chris@87: array([2, 1]) Chris@87: Chris@87: """ Chris@87: m = asanyarray(m) Chris@87: if m.ndim < 1: Chris@87: raise ValueError("Input must be >= 1-d.") Chris@87: return m[::-1, ...] Chris@87: Chris@87: Chris@87: def rot90(m, k=1): Chris@87: """ Chris@87: Rotate an array by 90 degrees in the counter-clockwise direction. Chris@87: Chris@87: The first two dimensions are rotated; therefore, the array must be at Chris@87: least 2-D. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: m : array_like Chris@87: Array of two or more dimensions. Chris@87: k : integer Chris@87: Number of times the array is rotated by 90 degrees. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: y : ndarray Chris@87: Rotated array. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: fliplr : Flip an array horizontally. Chris@87: flipud : Flip an array vertically. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> m = np.array([[1,2],[3,4]], int) Chris@87: >>> m Chris@87: array([[1, 2], Chris@87: [3, 4]]) Chris@87: >>> np.rot90(m) Chris@87: array([[2, 4], Chris@87: [1, 3]]) Chris@87: >>> np.rot90(m, 2) Chris@87: array([[4, 3], Chris@87: [2, 1]]) Chris@87: Chris@87: """ Chris@87: m = asanyarray(m) Chris@87: if m.ndim < 2: Chris@87: raise ValueError("Input must >= 2-d.") Chris@87: k = k % 4 Chris@87: if k == 0: Chris@87: return m Chris@87: elif k == 1: Chris@87: return fliplr(m).swapaxes(0, 1) Chris@87: elif k == 2: Chris@87: return fliplr(flipud(m)) Chris@87: else: Chris@87: # k == 3 Chris@87: return fliplr(m.swapaxes(0, 1)) Chris@87: Chris@87: Chris@87: def eye(N, M=None, k=0, dtype=float): Chris@87: """ Chris@87: Return a 2-D array with ones on the diagonal and zeros elsewhere. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: N : int Chris@87: Number of rows in the output. Chris@87: M : int, optional Chris@87: Number of columns in the output. If None, defaults to `N`. Chris@87: k : int, optional Chris@87: Index of the diagonal: 0 (the default) refers to the main diagonal, Chris@87: a positive value refers to an upper diagonal, and a negative value Chris@87: to a lower diagonal. Chris@87: dtype : data-type, optional Chris@87: Data-type of the returned array. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: I : ndarray of shape (N,M) Chris@87: An array where all elements are equal to zero, except for the `k`-th Chris@87: diagonal, whose values are equal to one. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: identity : (almost) equivalent function Chris@87: diag : diagonal 2-D array from a 1-D array specified by the user. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> np.eye(2, dtype=int) Chris@87: array([[1, 0], Chris@87: [0, 1]]) Chris@87: >>> np.eye(3, k=1) Chris@87: array([[ 0., 1., 0.], Chris@87: [ 0., 0., 1.], Chris@87: [ 0., 0., 0.]]) Chris@87: Chris@87: """ Chris@87: if M is None: Chris@87: M = N Chris@87: m = zeros((N, M), dtype=dtype) Chris@87: if k >= M: Chris@87: return m Chris@87: if k >= 0: Chris@87: i = k Chris@87: else: Chris@87: i = (-k) * M Chris@87: m[:M-k].flat[i::M+1] = 1 Chris@87: return m Chris@87: Chris@87: Chris@87: def diag(v, k=0): Chris@87: """ Chris@87: Extract a diagonal or construct a diagonal array. Chris@87: Chris@87: See the more detailed documentation for ``numpy.diagonal`` if you use this Chris@87: function to extract a diagonal and wish to write to the resulting array; Chris@87: whether it returns a copy or a view depends on what version of numpy you Chris@87: are using. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: v : array_like Chris@87: If `v` is a 2-D array, return a copy of its `k`-th diagonal. Chris@87: If `v` is a 1-D array, return a 2-D array with `v` on the `k`-th Chris@87: diagonal. Chris@87: k : int, optional Chris@87: Diagonal in question. The default is 0. Use `k>0` for diagonals Chris@87: above the main diagonal, and `k<0` for diagonals below the main Chris@87: diagonal. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: out : ndarray Chris@87: The extracted diagonal or constructed diagonal array. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: diagonal : Return specified diagonals. Chris@87: diagflat : Create a 2-D array with the flattened input as a diagonal. Chris@87: trace : Sum along diagonals. Chris@87: triu : Upper triangle of an array. Chris@87: tril : Lower triangle of an array. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> x = np.arange(9).reshape((3,3)) Chris@87: >>> x Chris@87: array([[0, 1, 2], Chris@87: [3, 4, 5], Chris@87: [6, 7, 8]]) Chris@87: Chris@87: >>> np.diag(x) Chris@87: array([0, 4, 8]) Chris@87: >>> np.diag(x, k=1) Chris@87: array([1, 5]) Chris@87: >>> np.diag(x, k=-1) Chris@87: array([3, 7]) Chris@87: Chris@87: >>> np.diag(np.diag(x)) Chris@87: array([[0, 0, 0], Chris@87: [0, 4, 0], Chris@87: [0, 0, 8]]) Chris@87: Chris@87: """ Chris@87: v = asarray(v) Chris@87: s = v.shape Chris@87: if len(s) == 1: Chris@87: n = s[0]+abs(k) Chris@87: res = zeros((n, n), v.dtype) Chris@87: if k >= 0: Chris@87: i = k Chris@87: else: Chris@87: i = (-k) * n Chris@87: res[:n-k].flat[i::n+1] = v Chris@87: return res Chris@87: elif len(s) == 2: Chris@87: return v.diagonal(k) Chris@87: else: Chris@87: raise ValueError("Input must be 1- or 2-d.") Chris@87: Chris@87: Chris@87: def diagflat(v, k=0): Chris@87: """ Chris@87: Create a two-dimensional array with the flattened input as a diagonal. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: v : array_like Chris@87: Input data, which is flattened and set as the `k`-th Chris@87: diagonal of the output. Chris@87: k : int, optional Chris@87: Diagonal to set; 0, the default, corresponds to the "main" diagonal, Chris@87: a positive (negative) `k` giving the number of the diagonal above Chris@87: (below) the main. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: out : ndarray Chris@87: The 2-D output array. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: diag : MATLAB work-alike for 1-D and 2-D arrays. Chris@87: diagonal : Return specified diagonals. Chris@87: trace : Sum along diagonals. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> np.diagflat([[1,2], [3,4]]) Chris@87: array([[1, 0, 0, 0], Chris@87: [0, 2, 0, 0], Chris@87: [0, 0, 3, 0], Chris@87: [0, 0, 0, 4]]) Chris@87: Chris@87: >>> np.diagflat([1,2], 1) Chris@87: array([[0, 1, 0], Chris@87: [0, 0, 2], Chris@87: [0, 0, 0]]) Chris@87: Chris@87: """ Chris@87: try: Chris@87: wrap = v.__array_wrap__ Chris@87: except AttributeError: Chris@87: wrap = None Chris@87: v = asarray(v).ravel() Chris@87: s = len(v) Chris@87: n = s + abs(k) Chris@87: res = zeros((n, n), v.dtype) Chris@87: if (k >= 0): Chris@87: i = arange(0, n-k) Chris@87: fi = i+k+i*n Chris@87: else: Chris@87: i = arange(0, n+k) Chris@87: fi = i+(i-k)*n Chris@87: res.flat[fi] = v Chris@87: if not wrap: Chris@87: return res Chris@87: return wrap(res) Chris@87: Chris@87: Chris@87: def tri(N, M=None, k=0, dtype=float): Chris@87: """ Chris@87: An array with ones at and below the given diagonal and zeros elsewhere. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: N : int Chris@87: Number of rows in the array. Chris@87: M : int, optional Chris@87: Number of columns in the array. Chris@87: By default, `M` is taken equal to `N`. Chris@87: k : int, optional Chris@87: The sub-diagonal at and below which the array is filled. Chris@87: `k` = 0 is the main diagonal, while `k` < 0 is below it, Chris@87: and `k` > 0 is above. The default is 0. Chris@87: dtype : dtype, optional Chris@87: Data type of the returned array. The default is float. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: tri : ndarray of shape (N, M) Chris@87: Array with its lower triangle filled with ones and zero elsewhere; Chris@87: in other words ``T[i,j] == 1`` for ``i <= j + k``, 0 otherwise. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> np.tri(3, 5, 2, dtype=int) Chris@87: array([[1, 1, 1, 0, 0], Chris@87: [1, 1, 1, 1, 0], Chris@87: [1, 1, 1, 1, 1]]) Chris@87: Chris@87: >>> np.tri(3, 5, -1) Chris@87: array([[ 0., 0., 0., 0., 0.], Chris@87: [ 1., 0., 0., 0., 0.], Chris@87: [ 1., 1., 0., 0., 0.]]) Chris@87: Chris@87: """ Chris@87: if M is None: Chris@87: M = N Chris@87: Chris@87: m = greater_equal.outer(arange(N, dtype=_min_int(0, N)), Chris@87: arange(-k, M-k, dtype=_min_int(-k, M - k))) Chris@87: Chris@87: # Avoid making a copy if the requested type is already bool Chris@87: m = m.astype(dtype, copy=False) Chris@87: Chris@87: return m Chris@87: Chris@87: Chris@87: def tril(m, k=0): Chris@87: """ Chris@87: Lower triangle of an array. Chris@87: Chris@87: Return a copy of an array with elements above the `k`-th diagonal zeroed. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: m : array_like, shape (M, N) Chris@87: Input array. Chris@87: k : int, optional Chris@87: Diagonal above which to zero elements. `k = 0` (the default) is the Chris@87: main diagonal, `k < 0` is below it and `k > 0` is above. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: tril : ndarray, shape (M, N) Chris@87: Lower triangle of `m`, of same shape and data-type as `m`. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: triu : same thing, only for the upper triangle Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> np.tril([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1) Chris@87: array([[ 0, 0, 0], Chris@87: [ 4, 0, 0], Chris@87: [ 7, 8, 0], Chris@87: [10, 11, 12]]) Chris@87: Chris@87: """ Chris@87: m = asanyarray(m) Chris@87: mask = tri(*m.shape[-2:], k=k, dtype=bool) Chris@87: Chris@87: return where(mask, m, zeros(1, m.dtype)) Chris@87: Chris@87: Chris@87: def triu(m, k=0): Chris@87: """ Chris@87: Upper triangle of an array. Chris@87: Chris@87: Return a copy of a matrix with the elements below the `k`-th diagonal Chris@87: zeroed. Chris@87: Chris@87: Please refer to the documentation for `tril` for further details. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: tril : lower triangle of an array Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1) Chris@87: array([[ 1, 2, 3], Chris@87: [ 4, 5, 6], Chris@87: [ 0, 8, 9], Chris@87: [ 0, 0, 12]]) Chris@87: Chris@87: """ Chris@87: m = asanyarray(m) Chris@87: mask = tri(*m.shape[-2:], k=k-1, dtype=bool) Chris@87: Chris@87: return where(mask, zeros(1, m.dtype), m) Chris@87: Chris@87: Chris@87: # Originally borrowed from John Hunter and matplotlib Chris@87: def vander(x, N=None, increasing=False): Chris@87: """ Chris@87: Generate a Vandermonde matrix. Chris@87: Chris@87: The columns of the output matrix are powers of the input vector. The Chris@87: order of the powers is determined by the `increasing` boolean argument. Chris@87: Specifically, when `increasing` is False, the `i`-th output column is Chris@87: the input vector raised element-wise to the power of ``N - i - 1``. Such Chris@87: a matrix with a geometric progression in each row is named for Alexandre- Chris@87: Theophile Vandermonde. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: x : array_like Chris@87: 1-D input array. Chris@87: N : int, optional Chris@87: Number of columns in the output. If `N` is not specified, a square Chris@87: array is returned (``N = len(x)``). Chris@87: increasing : bool, optional Chris@87: Order of the powers of the columns. If True, the powers increase Chris@87: from left to right, if False (the default) they are reversed. Chris@87: Chris@87: .. versionadded:: 1.9.0 Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: out : ndarray Chris@87: Vandermonde matrix. If `increasing` is False, the first column is Chris@87: ``x^(N-1)``, the second ``x^(N-2)`` and so forth. If `increasing` is Chris@87: True, the columns are ``x^0, x^1, ..., x^(N-1)``. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: polynomial.polynomial.polyvander Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> x = np.array([1, 2, 3, 5]) Chris@87: >>> N = 3 Chris@87: >>> np.vander(x, N) Chris@87: array([[ 1, 1, 1], Chris@87: [ 4, 2, 1], Chris@87: [ 9, 3, 1], Chris@87: [25, 5, 1]]) Chris@87: Chris@87: >>> np.column_stack([x**(N-1-i) for i in range(N)]) Chris@87: array([[ 1, 1, 1], Chris@87: [ 4, 2, 1], Chris@87: [ 9, 3, 1], Chris@87: [25, 5, 1]]) Chris@87: Chris@87: >>> x = np.array([1, 2, 3, 5]) Chris@87: >>> np.vander(x) Chris@87: array([[ 1, 1, 1, 1], Chris@87: [ 8, 4, 2, 1], Chris@87: [ 27, 9, 3, 1], Chris@87: [125, 25, 5, 1]]) Chris@87: >>> np.vander(x, increasing=True) Chris@87: array([[ 1, 1, 1, 1], Chris@87: [ 1, 2, 4, 8], Chris@87: [ 1, 3, 9, 27], Chris@87: [ 1, 5, 25, 125]]) Chris@87: Chris@87: The determinant of a square Vandermonde matrix is the product Chris@87: of the differences between the values of the input vector: Chris@87: Chris@87: >>> np.linalg.det(np.vander(x)) Chris@87: 48.000000000000043 Chris@87: >>> (5-3)*(5-2)*(5-1)*(3-2)*(3-1)*(2-1) Chris@87: 48 Chris@87: Chris@87: """ Chris@87: x = asarray(x) Chris@87: if x.ndim != 1: Chris@87: raise ValueError("x must be a one-dimensional array or sequence.") Chris@87: if N is None: Chris@87: N = len(x) Chris@87: Chris@87: v = empty((len(x), N), dtype=promote_types(x.dtype, int)) Chris@87: tmp = v[:, ::-1] if not increasing else v Chris@87: Chris@87: if N > 0: Chris@87: tmp[:, 0] = 1 Chris@87: if N > 1: Chris@87: tmp[:, 1:] = x[:, None] Chris@87: multiply.accumulate(tmp[:, 1:], out=tmp[:, 1:], axis=1) Chris@87: Chris@87: return v Chris@87: Chris@87: Chris@87: def histogram2d(x, y, bins=10, range=None, normed=False, weights=None): Chris@87: """ Chris@87: Compute the bi-dimensional histogram of two data samples. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: x : array_like, shape (N,) Chris@87: An array containing the x coordinates of the points to be Chris@87: histogrammed. Chris@87: y : array_like, shape (N,) Chris@87: An array containing the y coordinates of the points to be Chris@87: histogrammed. Chris@87: bins : int or [int, int] or array_like or [array, array], optional Chris@87: The bin specification: Chris@87: Chris@87: * If int, the number of bins for the two dimensions (nx=ny=bins). Chris@87: * If [int, int], the number of bins in each dimension Chris@87: (nx, ny = bins). Chris@87: * If array_like, the bin edges for the two dimensions Chris@87: (x_edges=y_edges=bins). Chris@87: * If [array, array], the bin edges in each dimension Chris@87: (x_edges, y_edges = bins). Chris@87: Chris@87: range : array_like, shape(2,2), optional Chris@87: The leftmost and rightmost edges of the bins along each dimension Chris@87: (if not specified explicitly in the `bins` parameters): Chris@87: ``[[xmin, xmax], [ymin, ymax]]``. All values outside of this range Chris@87: will be considered outliers and not tallied in the histogram. Chris@87: normed : bool, optional Chris@87: If False, returns the number of samples in each bin. If True, Chris@87: returns the bin density ``bin_count / sample_count / bin_area``. Chris@87: weights : array_like, shape(N,), optional Chris@87: An array of values ``w_i`` weighing each sample ``(x_i, y_i)``. Chris@87: Weights are normalized to 1 if `normed` is True. If `normed` is Chris@87: False, the values of the returned histogram are equal to the sum of Chris@87: the weights belonging to the samples falling into each bin. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: H : ndarray, shape(nx, ny) Chris@87: The bi-dimensional histogram of samples `x` and `y`. Values in `x` Chris@87: are histogrammed along the first dimension and values in `y` are Chris@87: histogrammed along the second dimension. Chris@87: xedges : ndarray, shape(nx,) Chris@87: The bin edges along the first dimension. Chris@87: yedges : ndarray, shape(ny,) Chris@87: The bin edges along the second dimension. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: histogram : 1D histogram Chris@87: histogramdd : Multidimensional histogram Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: When `normed` is True, then the returned histogram is the sample Chris@87: density, defined such that the sum over bins of the product Chris@87: ``bin_value * bin_area`` is 1. Chris@87: Chris@87: Please note that the histogram does not follow the Cartesian convention Chris@87: where `x` values are on the abscissa and `y` values on the ordinate Chris@87: axis. Rather, `x` is histogrammed along the first dimension of the Chris@87: array (vertical), and `y` along the second dimension of the array Chris@87: (horizontal). This ensures compatibility with `histogramdd`. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> import matplotlib as mpl Chris@87: >>> import matplotlib.pyplot as plt Chris@87: Chris@87: Construct a 2D-histogram with variable bin width. First define the bin Chris@87: edges: Chris@87: Chris@87: >>> xedges = [0, 1, 1.5, 3, 5] Chris@87: >>> yedges = [0, 2, 3, 4, 6] Chris@87: Chris@87: Next we create a histogram H with random bin content: Chris@87: Chris@87: >>> x = np.random.normal(3, 1, 100) Chris@87: >>> y = np.random.normal(1, 1, 100) Chris@87: >>> H, xedges, yedges = np.histogram2d(y, x, bins=(xedges, yedges)) Chris@87: Chris@87: Or we fill the histogram H with a determined bin content: Chris@87: Chris@87: >>> H = np.ones((4, 4)).cumsum().reshape(4, 4) Chris@87: >>> print H[::-1] # This shows the bin content in the order as plotted Chris@87: [[ 13. 14. 15. 16.] Chris@87: [ 9. 10. 11. 12.] Chris@87: [ 5. 6. 7. 8.] Chris@87: [ 1. 2. 3. 4.]] Chris@87: Chris@87: Imshow can only do an equidistant representation of bins: Chris@87: Chris@87: >>> fig = plt.figure(figsize=(7, 3)) Chris@87: >>> ax = fig.add_subplot(131) Chris@87: >>> ax.set_title('imshow: equidistant') Chris@87: >>> im = plt.imshow(H, interpolation='nearest', origin='low', Chris@87: extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]]) Chris@87: Chris@87: pcolormesh can display exact bin edges: Chris@87: Chris@87: >>> ax = fig.add_subplot(132) Chris@87: >>> ax.set_title('pcolormesh: exact bin edges') Chris@87: >>> X, Y = np.meshgrid(xedges, yedges) Chris@87: >>> ax.pcolormesh(X, Y, H) Chris@87: >>> ax.set_aspect('equal') Chris@87: Chris@87: NonUniformImage displays exact bin edges with interpolation: Chris@87: Chris@87: >>> ax = fig.add_subplot(133) Chris@87: >>> ax.set_title('NonUniformImage: interpolated') Chris@87: >>> im = mpl.image.NonUniformImage(ax, interpolation='bilinear') Chris@87: >>> xcenters = xedges[:-1] + 0.5 * (xedges[1:] - xedges[:-1]) Chris@87: >>> ycenters = yedges[:-1] + 0.5 * (yedges[1:] - yedges[:-1]) Chris@87: >>> im.set_data(xcenters, ycenters, H) Chris@87: >>> ax.images.append(im) Chris@87: >>> ax.set_xlim(xedges[0], xedges[-1]) Chris@87: >>> ax.set_ylim(yedges[0], yedges[-1]) Chris@87: >>> ax.set_aspect('equal') Chris@87: >>> plt.show() Chris@87: Chris@87: """ Chris@87: from numpy import histogramdd Chris@87: Chris@87: try: Chris@87: N = len(bins) Chris@87: except TypeError: Chris@87: N = 1 Chris@87: Chris@87: if N != 1 and N != 2: Chris@87: xedges = yedges = asarray(bins, float) Chris@87: bins = [xedges, yedges] Chris@87: hist, edges = histogramdd([x, y], bins, range, normed, weights) Chris@87: return hist, edges[0], edges[1] Chris@87: Chris@87: Chris@87: def mask_indices(n, mask_func, k=0): Chris@87: """ Chris@87: Return the indices to access (n, n) arrays, given a masking function. Chris@87: Chris@87: Assume `mask_func` is a function that, for a square array a of size Chris@87: ``(n, n)`` with a possible offset argument `k`, when called as Chris@87: ``mask_func(a, k)`` returns a new array with zeros in certain locations Chris@87: (functions like `triu` or `tril` do precisely this). Then this function Chris@87: returns the indices where the non-zero values would be located. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: n : int Chris@87: The returned indices will be valid to access arrays of shape (n, n). Chris@87: mask_func : callable Chris@87: A function whose call signature is similar to that of `triu`, `tril`. Chris@87: That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`. Chris@87: `k` is an optional argument to the function. Chris@87: k : scalar Chris@87: An optional argument which is passed through to `mask_func`. Functions Chris@87: like `triu`, `tril` take a second argument that is interpreted as an Chris@87: offset. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: indices : tuple of arrays. Chris@87: The `n` arrays of indices corresponding to the locations where Chris@87: ``mask_func(np.ones((n, n)), k)`` is True. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: triu, tril, triu_indices, tril_indices Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: .. versionadded:: 1.4.0 Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: These are the indices that would allow you to access the upper triangular Chris@87: part of any 3x3 array: Chris@87: Chris@87: >>> iu = np.mask_indices(3, np.triu) Chris@87: Chris@87: For example, if `a` is a 3x3 array: Chris@87: Chris@87: >>> a = np.arange(9).reshape(3, 3) Chris@87: >>> a Chris@87: array([[0, 1, 2], Chris@87: [3, 4, 5], Chris@87: [6, 7, 8]]) Chris@87: >>> a[iu] Chris@87: array([0, 1, 2, 4, 5, 8]) Chris@87: Chris@87: An offset can be passed also to the masking function. This gets us the Chris@87: indices starting on the first diagonal right of the main one: Chris@87: Chris@87: >>> iu1 = np.mask_indices(3, np.triu, 1) Chris@87: Chris@87: with which we now extract only three elements: Chris@87: Chris@87: >>> a[iu1] Chris@87: array([1, 2, 5]) Chris@87: Chris@87: """ Chris@87: m = ones((n, n), int) Chris@87: a = mask_func(m, k) Chris@87: return where(a != 0) Chris@87: Chris@87: Chris@87: def tril_indices(n, k=0, m=None): Chris@87: """ Chris@87: Return the indices for the lower-triangle of an (n, m) array. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: n : int Chris@87: The row dimension of the arrays for which the returned Chris@87: indices will be valid. Chris@87: k : int, optional Chris@87: Diagonal offset (see `tril` for details). Chris@87: m : int, optional Chris@87: .. versionadded:: 1.9.0 Chris@87: Chris@87: The column dimension of the arrays for which the returned Chris@87: arrays will be valid. Chris@87: By default `m` is taken equal to `n`. Chris@87: Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: inds : tuple of arrays Chris@87: The indices for the triangle. The returned tuple contains two arrays, Chris@87: each with the indices along one dimension of the array. Chris@87: Chris@87: See also Chris@87: -------- Chris@87: triu_indices : similar function, for upper-triangular. Chris@87: mask_indices : generic function accepting an arbitrary mask function. Chris@87: tril, triu Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: .. versionadded:: 1.4.0 Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: Compute two different sets of indices to access 4x4 arrays, one for the Chris@87: lower triangular part starting at the main diagonal, and one starting two Chris@87: diagonals further right: Chris@87: Chris@87: >>> il1 = np.tril_indices(4) Chris@87: >>> il2 = np.tril_indices(4, 2) Chris@87: Chris@87: Here is how they can be used with a sample array: Chris@87: Chris@87: >>> a = np.arange(16).reshape(4, 4) Chris@87: >>> a Chris@87: array([[ 0, 1, 2, 3], Chris@87: [ 4, 5, 6, 7], Chris@87: [ 8, 9, 10, 11], Chris@87: [12, 13, 14, 15]]) Chris@87: Chris@87: Both for indexing: Chris@87: Chris@87: >>> a[il1] Chris@87: array([ 0, 4, 5, 8, 9, 10, 12, 13, 14, 15]) Chris@87: Chris@87: And for assigning values: Chris@87: Chris@87: >>> a[il1] = -1 Chris@87: >>> a Chris@87: array([[-1, 1, 2, 3], Chris@87: [-1, -1, 6, 7], Chris@87: [-1, -1, -1, 11], Chris@87: [-1, -1, -1, -1]]) Chris@87: Chris@87: These cover almost the whole array (two diagonals right of the main one): Chris@87: Chris@87: >>> a[il2] = -10 Chris@87: >>> a Chris@87: array([[-10, -10, -10, 3], Chris@87: [-10, -10, -10, -10], Chris@87: [-10, -10, -10, -10], Chris@87: [-10, -10, -10, -10]]) Chris@87: Chris@87: """ Chris@87: return where(tri(n, m, k=k, dtype=bool)) Chris@87: Chris@87: Chris@87: def tril_indices_from(arr, k=0): Chris@87: """ Chris@87: Return the indices for the lower-triangle of arr. Chris@87: Chris@87: See `tril_indices` for full details. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: arr : array_like Chris@87: The indices will be valid for square arrays whose dimensions are Chris@87: the same as arr. Chris@87: k : int, optional Chris@87: Diagonal offset (see `tril` for details). Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: tril_indices, tril Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: .. versionadded:: 1.4.0 Chris@87: Chris@87: """ Chris@87: if arr.ndim != 2: Chris@87: raise ValueError("input array must be 2-d") Chris@87: return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1]) Chris@87: Chris@87: Chris@87: def triu_indices(n, k=0, m=None): Chris@87: """ Chris@87: Return the indices for the upper-triangle of an (n, m) array. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: n : int Chris@87: The size of the arrays for which the returned indices will Chris@87: be valid. Chris@87: k : int, optional Chris@87: Diagonal offset (see `triu` for details). Chris@87: m : int, optional Chris@87: .. versionadded:: 1.9.0 Chris@87: Chris@87: The column dimension of the arrays for which the returned Chris@87: arrays will be valid. Chris@87: By default `m` is taken equal to `n`. Chris@87: Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: inds : tuple, shape(2) of ndarrays, shape(`n`) Chris@87: The indices for the triangle. The returned tuple contains two arrays, Chris@87: each with the indices along one dimension of the array. Can be used Chris@87: to slice a ndarray of shape(`n`, `n`). Chris@87: Chris@87: See also Chris@87: -------- Chris@87: tril_indices : similar function, for lower-triangular. Chris@87: mask_indices : generic function accepting an arbitrary mask function. Chris@87: triu, tril Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: .. versionadded:: 1.4.0 Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: Compute two different sets of indices to access 4x4 arrays, one for the Chris@87: upper triangular part starting at the main diagonal, and one starting two Chris@87: diagonals further right: Chris@87: Chris@87: >>> iu1 = np.triu_indices(4) Chris@87: >>> iu2 = np.triu_indices(4, 2) Chris@87: Chris@87: Here is how they can be used with a sample array: Chris@87: Chris@87: >>> a = np.arange(16).reshape(4, 4) Chris@87: >>> a Chris@87: array([[ 0, 1, 2, 3], Chris@87: [ 4, 5, 6, 7], Chris@87: [ 8, 9, 10, 11], Chris@87: [12, 13, 14, 15]]) Chris@87: Chris@87: Both for indexing: Chris@87: Chris@87: >>> a[iu1] Chris@87: array([ 0, 1, 2, 3, 5, 6, 7, 10, 11, 15]) Chris@87: Chris@87: And for assigning values: Chris@87: Chris@87: >>> a[iu1] = -1 Chris@87: >>> a Chris@87: array([[-1, -1, -1, -1], Chris@87: [ 4, -1, -1, -1], Chris@87: [ 8, 9, -1, -1], Chris@87: [12, 13, 14, -1]]) Chris@87: Chris@87: These cover only a small part of the whole array (two diagonals right Chris@87: of the main one): Chris@87: Chris@87: >>> a[iu2] = -10 Chris@87: >>> a Chris@87: array([[ -1, -1, -10, -10], Chris@87: [ 4, -1, -1, -10], Chris@87: [ 8, 9, -1, -1], Chris@87: [ 12, 13, 14, -1]]) Chris@87: Chris@87: """ Chris@87: return where(~tri(n, m, k=k-1, dtype=bool)) Chris@87: Chris@87: Chris@87: def triu_indices_from(arr, k=0): Chris@87: """ Chris@87: Return the indices for the upper-triangle of arr. Chris@87: Chris@87: See `triu_indices` for full details. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: arr : ndarray, shape(N, N) Chris@87: The indices will be valid for square arrays. Chris@87: k : int, optional Chris@87: Diagonal offset (see `triu` for details). Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: triu_indices_from : tuple, shape(2) of ndarray, shape(N) Chris@87: Indices for the upper-triangle of `arr`. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: triu_indices, triu Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: .. versionadded:: 1.4.0 Chris@87: Chris@87: """ Chris@87: if arr.ndim != 2: Chris@87: raise ValueError("input array must be 2-d") Chris@87: return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1])