MLPY/Lib/site-packages/numpy/lib/twodim_base.py

""" Basic functions for manipulating 2d arrays



"""
import functools

from numpy.core.numeric import (
    asanyarray, arange, zeros, greater_equal, multiply, ones,
    asarray, where, int8, int16, int32, int64, intp, empty, promote_types,
    diagonal, nonzero, indices
    )
from numpy.core.overrides import set_array_function_like_doc, set_module
from numpy.core import overrides
from numpy.core import iinfo
from numpy.lib.stride_tricks import broadcast_to


__all__ = [
    'diag', 'diagflat', 'eye', 'fliplr', 'flipud', 'tri', 'triu',
    'tril', 'vander', 'histogram2d', 'mask_indices', 'tril_indices',
    'tril_indices_from', 'triu_indices', 'triu_indices_from', ]


array_function_dispatch = functools.partial(
    overrides.array_function_dispatch, module='numpy')


i1 = iinfo(int8)
i2 = iinfo(int16)
i4 = iinfo(int32)


def _min_int(low, high):
    """ get small int that fits the range """
    if high <= i1.max and low >= i1.min:
        return int8
    if high <= i2.max and low >= i2.min:
        return int16
    if high <= i4.max and low >= i4.min:
        return int32
    return int64


def _flip_dispatcher(m):
    return (m,)


@array_function_dispatch(_flip_dispatcher)
def fliplr(m):
    """

    Reverse the order of elements along axis 1 (left/right).



    For a 2-D array, this flips the entries in each row in the left/right

    direction. Columns are preserved, but appear in a different order than

    before.



    Parameters

    ----------

    m : array_like

        Input array, must be at least 2-D.



    Returns

    -------

    f : ndarray

        A view of `m` with the columns reversed.  Since a view

        is returned, this operation is :math:`\\mathcal O(1)`.



    See Also

    --------

    flipud : Flip array in the up/down direction.

    flip : Flip array in one or more dimesions.

    rot90 : Rotate array counterclockwise.



    Notes

    -----

    Equivalent to ``m[:,::-1]`` or ``np.flip(m, axis=1)``.

    Requires the array to be at least 2-D.



    Examples

    --------

    >>> A = np.diag([1.,2.,3.])

    >>> A

    array([[1.,  0.,  0.],

           [0.,  2.,  0.],

           [0.,  0.,  3.]])

    >>> np.fliplr(A)

    array([[0.,  0.,  1.],

           [0.,  2.,  0.],

           [3.,  0.,  0.]])



    >>> A = np.random.randn(2,3,5)

    >>> np.all(np.fliplr(A) == A[:,::-1,...])

    True



    """
    m = asanyarray(m)
    if m.ndim < 2:
        raise ValueError("Input must be >= 2-d.")
    return m[:, ::-1]


@array_function_dispatch(_flip_dispatcher)
def flipud(m):
    """

    Reverse the order of elements along axis 0 (up/down).



    For a 2-D array, this flips the entries in each column in the up/down

    direction. Rows are preserved, but appear in a different order than before.



    Parameters

    ----------

    m : array_like

        Input array.



    Returns

    -------

    out : array_like

        A view of `m` with the rows reversed.  Since a view is

        returned, this operation is :math:`\\mathcal O(1)`.



    See Also

    --------

    fliplr : Flip array in the left/right direction.

    flip : Flip array in one or more dimesions.

    rot90 : Rotate array counterclockwise.



    Notes

    -----

    Equivalent to ``m[::-1, ...]`` or ``np.flip(m, axis=0)``.

    Requires the array to be at least 1-D.



    Examples

    --------

    >>> A = np.diag([1.0, 2, 3])

    >>> A

    array([[1.,  0.,  0.],

           [0.,  2.,  0.],

           [0.,  0.,  3.]])

    >>> np.flipud(A)

    array([[0.,  0.,  3.],

           [0.,  2.,  0.],

           [1.,  0.,  0.]])



    >>> A = np.random.randn(2,3,5)

    >>> np.all(np.flipud(A) == A[::-1,...])

    True



    >>> np.flipud([1,2])

    array([2, 1])



    """
    m = asanyarray(m)
    if m.ndim < 1:
        raise ValueError("Input must be >= 1-d.")
    return m[::-1, ...]


def _eye_dispatcher(N, M=None, k=None, dtype=None, order=None, *, like=None):
    return (like,)


@set_array_function_like_doc
@set_module('numpy')
def eye(N, M=None, k=0, dtype=float, order='C', *, like=None):
    """

    Return a 2-D array with ones on the diagonal and zeros elsewhere.



    Parameters

    ----------

    N : int

      Number of rows in the output.

    M : int, optional

      Number of columns in the output. If None, defaults to `N`.

    k : int, optional

      Index of the diagonal: 0 (the default) refers to the main diagonal,

      a positive value refers to an upper diagonal, and a negative value

      to a lower diagonal.

    dtype : data-type, optional

      Data-type of the returned array.

    order : {'C', 'F'}, optional

        Whether the output should be stored in row-major (C-style) or

        column-major (Fortran-style) order in memory.



        .. versionadded:: 1.14.0

    ${ARRAY_FUNCTION_LIKE}



        .. versionadded:: 1.20.0



    Returns

    -------

    I : ndarray of shape (N,M)

      An array where all elements are equal to zero, except for the `k`-th

      diagonal, whose values are equal to one.



    See Also

    --------

    identity : (almost) equivalent function

    diag : diagonal 2-D array from a 1-D array specified by the user.



    Examples

    --------

    >>> np.eye(2, dtype=int)

    array([[1, 0],

           [0, 1]])

    >>> np.eye(3, k=1)

    array([[0.,  1.,  0.],

           [0.,  0.,  1.],

           [0.,  0.,  0.]])



    """
    if like is not None:
        return _eye_with_like(N, M=M, k=k, dtype=dtype, order=order, like=like)
    if M is None:
        M = N
    m = zeros((N, M), dtype=dtype, order=order)
    if k >= M:
        return m
    if k >= 0:
        i = k
    else:
        i = (-k) * M
    m[:M-k].flat[i::M+1] = 1
    return m


_eye_with_like = array_function_dispatch(
    _eye_dispatcher
)(eye)


def _diag_dispatcher(v, k=None):
    return (v,)


@array_function_dispatch(_diag_dispatcher)
def diag(v, k=0):
    """

    Extract a diagonal or construct a diagonal array.



    See the more detailed documentation for ``numpy.diagonal`` if you use this

    function to extract a diagonal and wish to write to the resulting array;

    whether it returns a copy or a view depends on what version of numpy you

    are using.



    Parameters

    ----------

    v : array_like

        If `v` is a 2-D array, return a copy of its `k`-th diagonal.

        If `v` is a 1-D array, return a 2-D array with `v` on the `k`-th

        diagonal.

    k : int, optional

        Diagonal in question. The default is 0. Use `k>0` for diagonals

        above the main diagonal, and `k<0` for diagonals below the main

        diagonal.



    Returns

    -------

    out : ndarray

        The extracted diagonal or constructed diagonal array.



    See Also

    --------

    diagonal : Return specified diagonals.

    diagflat : Create a 2-D array with the flattened input as a diagonal.

    trace : Sum along diagonals.

    triu : Upper triangle of an array.

    tril : Lower triangle of an array.



    Examples

    --------

    >>> x = np.arange(9).reshape((3,3))

    >>> x

    array([[0, 1, 2],

           [3, 4, 5],

           [6, 7, 8]])



    >>> np.diag(x)

    array([0, 4, 8])

    >>> np.diag(x, k=1)

    array([1, 5])

    >>> np.diag(x, k=-1)

    array([3, 7])



    >>> np.diag(np.diag(x))

    array([[0, 0, 0],

           [0, 4, 0],

           [0, 0, 8]])



    """
    v = asanyarray(v)
    s = v.shape
    if len(s) == 1:
        n = s[0]+abs(k)
        res = zeros((n, n), v.dtype)
        if k >= 0:
            i = k
        else:
            i = (-k) * n
        res[:n-k].flat[i::n+1] = v
        return res
    elif len(s) == 2:
        return diagonal(v, k)
    else:
        raise ValueError("Input must be 1- or 2-d.")


@array_function_dispatch(_diag_dispatcher)
def diagflat(v, k=0):
    """

    Create a two-dimensional array with the flattened input as a diagonal.



    Parameters

    ----------

    v : array_like

        Input data, which is flattened and set as the `k`-th

        diagonal of the output.

    k : int, optional

        Diagonal to set; 0, the default, corresponds to the "main" diagonal,

        a positive (negative) `k` giving the number of the diagonal above

        (below) the main.



    Returns

    -------

    out : ndarray

        The 2-D output array.



    See Also

    --------

    diag : MATLAB work-alike for 1-D and 2-D arrays.

    diagonal : Return specified diagonals.

    trace : Sum along diagonals.



    Examples

    --------

    >>> np.diagflat([[1,2], [3,4]])

    array([[1, 0, 0, 0],

           [0, 2, 0, 0],

           [0, 0, 3, 0],

           [0, 0, 0, 4]])



    >>> np.diagflat([1,2], 1)

    array([[0, 1, 0],

           [0, 0, 2],

           [0, 0, 0]])



    """
    try:
        wrap = v.__array_wrap__
    except AttributeError:
        wrap = None
    v = asarray(v).ravel()
    s = len(v)
    n = s + abs(k)
    res = zeros((n, n), v.dtype)
    if (k >= 0):
        i = arange(0, n-k, dtype=intp)
        fi = i+k+i*n
    else:
        i = arange(0, n+k, dtype=intp)
        fi = i+(i-k)*n
    res.flat[fi] = v
    if not wrap:
        return res
    return wrap(res)


def _tri_dispatcher(N, M=None, k=None, dtype=None, *, like=None):
    return (like,)


@set_array_function_like_doc
@set_module('numpy')
def tri(N, M=None, k=0, dtype=float, *, like=None):
    """

    An array with ones at and below the given diagonal and zeros elsewhere.



    Parameters

    ----------

    N : int

        Number of rows in the array.

    M : int, optional

        Number of columns in the array.

        By default, `M` is taken equal to `N`.

    k : int, optional

        The sub-diagonal at and below which the array is filled.

        `k` = 0 is the main diagonal, while `k` < 0 is below it,

        and `k` > 0 is above.  The default is 0.

    dtype : dtype, optional

        Data type of the returned array.  The default is float.

    ${ARRAY_FUNCTION_LIKE}



        .. versionadded:: 1.20.0



    Returns

    -------

    tri : ndarray of shape (N, M)

        Array with its lower triangle filled with ones and zero elsewhere;

        in other words ``T[i,j] == 1`` for ``j <= i + k``, 0 otherwise.



    Examples

    --------

    >>> np.tri(3, 5, 2, dtype=int)

    array([[1, 1, 1, 0, 0],

           [1, 1, 1, 1, 0],

           [1, 1, 1, 1, 1]])



    >>> np.tri(3, 5, -1)

    array([[0.,  0.,  0.,  0.,  0.],

           [1.,  0.,  0.,  0.,  0.],

           [1.,  1.,  0.,  0.,  0.]])



    """
    if like is not None:
        return _tri_with_like(N, M=M, k=k, dtype=dtype, like=like)

    if M is None:
        M = N

    m = greater_equal.outer(arange(N, dtype=_min_int(0, N)),
                            arange(-k, M-k, dtype=_min_int(-k, M - k)))

    # Avoid making a copy if the requested type is already bool
    m = m.astype(dtype, copy=False)

    return m


_tri_with_like = array_function_dispatch(
    _tri_dispatcher
)(tri)


def _trilu_dispatcher(m, k=None):
    return (m,)


@array_function_dispatch(_trilu_dispatcher)
def tril(m, k=0):
    """

    Lower triangle of an array.



    Return a copy of an array with elements above the `k`-th diagonal zeroed.



    Parameters

    ----------

    m : array_like, shape (M, N)

        Input array.

    k : int, optional

        Diagonal above which to zero elements.  `k = 0` (the default) is the

        main diagonal, `k < 0` is below it and `k > 0` is above.



    Returns

    -------

    tril : ndarray, shape (M, N)

        Lower triangle of `m`, of same shape and data-type as `m`.



    See Also

    --------

    triu : same thing, only for the upper triangle



    Examples

    --------

    >>> np.tril([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1)

    array([[ 0,  0,  0],

           [ 4,  0,  0],

           [ 7,  8,  0],

           [10, 11, 12]])



    """
    m = asanyarray(m)
    mask = tri(*m.shape[-2:], k=k, dtype=bool)

    return where(mask, m, zeros(1, m.dtype))


@array_function_dispatch(_trilu_dispatcher)
def triu(m, k=0):
    """

    Upper triangle of an array.



    Return a copy of an array with the elements below the `k`-th diagonal

    zeroed.



    Please refer to the documentation for `tril` for further details.



    See Also

    --------

    tril : lower triangle of an array



    Examples

    --------

    >>> np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1)

    array([[ 1,  2,  3],

           [ 4,  5,  6],

           [ 0,  8,  9],

           [ 0,  0, 12]])



    """
    m = asanyarray(m)
    mask = tri(*m.shape[-2:], k=k-1, dtype=bool)

    return where(mask, zeros(1, m.dtype), m)


def _vander_dispatcher(x, N=None, increasing=None):
    return (x,)


# Originally borrowed from John Hunter and matplotlib
@array_function_dispatch(_vander_dispatcher)
def vander(x, N=None, increasing=False):
    """

    Generate a Vandermonde matrix.



    The columns of the output matrix are powers of the input vector. The

    order of the powers is determined by the `increasing` boolean argument.

    Specifically, when `increasing` is False, the `i`-th output column is

    the input vector raised element-wise to the power of ``N - i - 1``. Such

    a matrix with a geometric progression in each row is named for Alexandre-

    Theophile Vandermonde.



    Parameters

    ----------

    x : array_like

        1-D input array.

    N : int, optional

        Number of columns in the output.  If `N` is not specified, a square

        array is returned (``N = len(x)``).

    increasing : bool, optional

        Order of the powers of the columns.  If True, the powers increase

        from left to right, if False (the default) they are reversed.



        .. versionadded:: 1.9.0



    Returns

    -------

    out : ndarray

        Vandermonde matrix.  If `increasing` is False, the first column is

        ``x^(N-1)``, the second ``x^(N-2)`` and so forth. If `increasing` is

        True, the columns are ``x^0, x^1, ..., x^(N-1)``.



    See Also

    --------

    polynomial.polynomial.polyvander



    Examples

    --------

    >>> x = np.array([1, 2, 3, 5])

    >>> N = 3

    >>> np.vander(x, N)

    array([[ 1,  1,  1],

           [ 4,  2,  1],

           [ 9,  3,  1],

           [25,  5,  1]])



    >>> np.column_stack([x**(N-1-i) for i in range(N)])

    array([[ 1,  1,  1],

           [ 4,  2,  1],

           [ 9,  3,  1],

           [25,  5,  1]])



    >>> x = np.array([1, 2, 3, 5])

    >>> np.vander(x)

    array([[  1,   1,   1,   1],

           [  8,   4,   2,   1],

           [ 27,   9,   3,   1],

           [125,  25,   5,   1]])

    >>> np.vander(x, increasing=True)

    array([[  1,   1,   1,   1],

           [  1,   2,   4,   8],

           [  1,   3,   9,  27],

           [  1,   5,  25, 125]])



    The determinant of a square Vandermonde matrix is the product

    of the differences between the values of the input vector:



    >>> np.linalg.det(np.vander(x))

    48.000000000000043 # may vary

    >>> (5-3)*(5-2)*(5-1)*(3-2)*(3-1)*(2-1)

    48



    """
    x = asarray(x)
    if x.ndim != 1:
        raise ValueError("x must be a one-dimensional array or sequence.")
    if N is None:
        N = len(x)

    v = empty((len(x), N), dtype=promote_types(x.dtype, int))
    tmp = v[:, ::-1] if not increasing else v

    if N > 0:
        tmp[:, 0] = 1
    if N > 1:
        tmp[:, 1:] = x[:, None]
        multiply.accumulate(tmp[:, 1:], out=tmp[:, 1:], axis=1)

    return v


def _histogram2d_dispatcher(x, y, bins=None, range=None, normed=None,

                            weights=None, density=None):
    yield x
    yield y

    # This terrible logic is adapted from the checks in histogram2d
    try:
        N = len(bins)
    except TypeError:
        N = 1
    if N == 2:
        yield from bins  # bins=[x, y]
    else:
        yield bins

    yield weights


@array_function_dispatch(_histogram2d_dispatcher)
def histogram2d(x, y, bins=10, range=None, normed=None, weights=None,

                density=None):
    """

    Compute the bi-dimensional histogram of two data samples.



    Parameters

    ----------

    x : array_like, shape (N,)

        An array containing the x coordinates of the points to be

        histogrammed.

    y : array_like, shape (N,)

        An array containing the y coordinates of the points to be

        histogrammed.

    bins : int or array_like or [int, int] or [array, array], optional

        The bin specification:



          * If int, the number of bins for the two dimensions (nx=ny=bins).

          * If array_like, the bin edges for the two dimensions

            (x_edges=y_edges=bins).

          * If [int, int], the number of bins in each dimension

            (nx, ny = bins).

          * If [array, array], the bin edges in each dimension

            (x_edges, y_edges = bins).

          * A combination [int, array] or [array, int], where int

            is the number of bins and array is the bin edges.



    range : array_like, shape(2,2), optional

        The leftmost and rightmost edges of the bins along each dimension

        (if not specified explicitly in the `bins` parameters):

        ``[[xmin, xmax], [ymin, ymax]]``. All values outside of this range

        will be considered outliers and not tallied in the histogram.

    density : bool, optional

        If False, the default, returns the number of samples in each bin.

        If True, returns the probability *density* function at the bin,

        ``bin_count / sample_count / bin_area``.

    normed : bool, optional

        An alias for the density argument that behaves identically. To avoid

        confusion with the broken normed argument to `histogram`, `density`

        should be preferred.

    weights : array_like, shape(N,), optional

        An array of values ``w_i`` weighing each sample ``(x_i, y_i)``.

        Weights are normalized to 1 if `normed` is True. If `normed` is

        False, the values of the returned histogram are equal to the sum of

        the weights belonging to the samples falling into each bin.



    Returns

    -------

    H : ndarray, shape(nx, ny)

        The bi-dimensional histogram of samples `x` and `y`. Values in `x`

        are histogrammed along the first dimension and values in `y` are

        histogrammed along the second dimension.

    xedges : ndarray, shape(nx+1,)

        The bin edges along the first dimension.

    yedges : ndarray, shape(ny+1,)

        The bin edges along the second dimension.



    See Also

    --------

    histogram : 1D histogram

    histogramdd : Multidimensional histogram



    Notes

    -----

    When `normed` is True, then the returned histogram is the sample

    density, defined such that the sum over bins of the product

    ``bin_value * bin_area`` is 1.



    Please note that the histogram does not follow the Cartesian convention

    where `x` values are on the abscissa and `y` values on the ordinate

    axis.  Rather, `x` is histogrammed along the first dimension of the

    array (vertical), and `y` along the second dimension of the array

    (horizontal).  This ensures compatibility with `histogramdd`.



    Examples

    --------

    >>> from matplotlib.image import NonUniformImage

    >>> import matplotlib.pyplot as plt



    Construct a 2-D histogram with variable bin width. First define the bin

    edges:



    >>> xedges = [0, 1, 3, 5]

    >>> yedges = [0, 2, 3, 4, 6]



    Next we create a histogram H with random bin content:



    >>> x = np.random.normal(2, 1, 100)

    >>> y = np.random.normal(1, 1, 100)

    >>> H, xedges, yedges = np.histogram2d(x, y, bins=(xedges, yedges))

    >>> # Histogram does not follow Cartesian convention (see Notes),

    >>> # therefore transpose H for visualization purposes.

    >>> H = H.T



    :func:`imshow <matplotlib.pyplot.imshow>` can only display square bins:



    >>> fig = plt.figure(figsize=(7, 3))

    >>> ax = fig.add_subplot(131, title='imshow: square bins')

    >>> plt.imshow(H, interpolation='nearest', origin='lower',

    ...         extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]])

    <matplotlib.image.AxesImage object at 0x...>



    :func:`pcolormesh <matplotlib.pyplot.pcolormesh>` can display actual edges:



    >>> ax = fig.add_subplot(132, title='pcolormesh: actual edges',

    ...         aspect='equal')

    >>> X, Y = np.meshgrid(xedges, yedges)

    >>> ax.pcolormesh(X, Y, H)

    <matplotlib.collections.QuadMesh object at 0x...>



    :class:`NonUniformImage <matplotlib.image.NonUniformImage>` can be used to

    display actual bin edges with interpolation:



    >>> ax = fig.add_subplot(133, title='NonUniformImage: interpolated',

    ...         aspect='equal', xlim=xedges[[0, -1]], ylim=yedges[[0, -1]])

    >>> im = NonUniformImage(ax, interpolation='bilinear')

    >>> xcenters = (xedges[:-1] + xedges[1:]) / 2

    >>> ycenters = (yedges[:-1] + yedges[1:]) / 2

    >>> im.set_data(xcenters, ycenters, H)

    >>> ax.images.append(im)

    >>> plt.show()



    """
    from numpy import histogramdd

    try:
        N = len(bins)
    except TypeError:
        N = 1

    if N != 1 and N != 2:
        xedges = yedges = asarray(bins)
        bins = [xedges, yedges]
    hist, edges = histogramdd([x, y], bins, range, normed, weights, density)
    return hist, edges[0], edges[1]


@set_module('numpy')
def mask_indices(n, mask_func, k=0):
    """

    Return the indices to access (n, n) arrays, given a masking function.



    Assume `mask_func` is a function that, for a square array a of size

    ``(n, n)`` with a possible offset argument `k`, when called as

    ``mask_func(a, k)`` returns a new array with zeros in certain locations

    (functions like `triu` or `tril` do precisely this). Then this function

    returns the indices where the non-zero values would be located.



    Parameters

    ----------

    n : int

        The returned indices will be valid to access arrays of shape (n, n).

    mask_func : callable

        A function whose call signature is similar to that of `triu`, `tril`.

        That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`.

        `k` is an optional argument to the function.

    k : scalar

        An optional argument which is passed through to `mask_func`. Functions

        like `triu`, `tril` take a second argument that is interpreted as an

        offset.



    Returns

    -------

    indices : tuple of arrays.

        The `n` arrays of indices corresponding to the locations where

        ``mask_func(np.ones((n, n)), k)`` is True.



    See Also

    --------

    triu, tril, triu_indices, tril_indices



    Notes

    -----

    .. versionadded:: 1.4.0



    Examples

    --------

    These are the indices that would allow you to access the upper triangular

    part of any 3x3 array:



    >>> iu = np.mask_indices(3, np.triu)



    For example, if `a` is a 3x3 array:



    >>> a = np.arange(9).reshape(3, 3)

    >>> a

    array([[0, 1, 2],

           [3, 4, 5],

           [6, 7, 8]])

    >>> a[iu]

    array([0, 1, 2, 4, 5, 8])



    An offset can be passed also to the masking function.  This gets us the

    indices starting on the first diagonal right of the main one:



    >>> iu1 = np.mask_indices(3, np.triu, 1)



    with which we now extract only three elements:



    >>> a[iu1]

    array([1, 2, 5])



    """
    m = ones((n, n), int)
    a = mask_func(m, k)
    return nonzero(a != 0)


@set_module('numpy')
def tril_indices(n, k=0, m=None):
    """

    Return the indices for the lower-triangle of an (n, m) array.



    Parameters

    ----------

    n : int

        The row dimension of the arrays for which the returned

        indices will be valid.

    k : int, optional

        Diagonal offset (see `tril` for details).

    m : int, optional

        .. versionadded:: 1.9.0



        The column dimension of the arrays for which the returned

        arrays will be valid.

        By default `m` is taken equal to `n`.





    Returns

    -------

    inds : tuple of arrays

        The indices for the triangle. The returned tuple contains two arrays,

        each with the indices along one dimension of the array.



    See also

    --------

    triu_indices : similar function, for upper-triangular.

    mask_indices : generic function accepting an arbitrary mask function.

    tril, triu



    Notes

    -----

    .. versionadded:: 1.4.0



    Examples

    --------

    Compute two different sets of indices to access 4x4 arrays, one for the

    lower triangular part starting at the main diagonal, and one starting two

    diagonals further right:



    >>> il1 = np.tril_indices(4)

    >>> il2 = np.tril_indices(4, 2)



    Here is how they can be used with a sample array:



    >>> a = np.arange(16).reshape(4, 4)

    >>> a

    array([[ 0,  1,  2,  3],

           [ 4,  5,  6,  7],

           [ 8,  9, 10, 11],

           [12, 13, 14, 15]])



    Both for indexing:



    >>> a[il1]

    array([ 0,  4,  5, ..., 13, 14, 15])



    And for assigning values:



    >>> a[il1] = -1

    >>> a

    array([[-1,  1,  2,  3],

           [-1, -1,  6,  7],

           [-1, -1, -1, 11],

           [-1, -1, -1, -1]])



    These cover almost the whole array (two diagonals right of the main one):



    >>> a[il2] = -10

    >>> a

    array([[-10, -10, -10,   3],

           [-10, -10, -10, -10],

           [-10, -10, -10, -10],

           [-10, -10, -10, -10]])



    """
    tri_ = tri(n, m, k=k, dtype=bool)

    return tuple(broadcast_to(inds, tri_.shape)[tri_]
                 for inds in indices(tri_.shape, sparse=True))


def _trilu_indices_form_dispatcher(arr, k=None):
    return (arr,)


@array_function_dispatch(_trilu_indices_form_dispatcher)
def tril_indices_from(arr, k=0):
    """

    Return the indices for the lower-triangle of arr.



    See `tril_indices` for full details.



    Parameters

    ----------

    arr : array_like

        The indices will be valid for square arrays whose dimensions are

        the same as arr.

    k : int, optional

        Diagonal offset (see `tril` for details).



    See Also

    --------

    tril_indices, tril



    Notes

    -----

    .. versionadded:: 1.4.0



    """
    if arr.ndim != 2:
        raise ValueError("input array must be 2-d")
    return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1])


@set_module('numpy')
def triu_indices(n, k=0, m=None):
    """

    Return the indices for the upper-triangle of an (n, m) array.



    Parameters

    ----------

    n : int

        The size of the arrays for which the returned indices will

        be valid.

    k : int, optional

        Diagonal offset (see `triu` for details).

    m : int, optional

        .. versionadded:: 1.9.0



        The column dimension of the arrays for which the returned

        arrays will be valid.

        By default `m` is taken equal to `n`.





    Returns

    -------

    inds : tuple, shape(2) of ndarrays, shape(`n`)

        The indices for the triangle. The returned tuple contains two arrays,

        each with the indices along one dimension of the array.  Can be used

        to slice a ndarray of shape(`n`, `n`).



    See also

    --------

    tril_indices : similar function, for lower-triangular.

    mask_indices : generic function accepting an arbitrary mask function.

    triu, tril



    Notes

    -----

    .. versionadded:: 1.4.0



    Examples

    --------

    Compute two different sets of indices to access 4x4 arrays, one for the

    upper triangular part starting at the main diagonal, and one starting two

    diagonals further right:



    >>> iu1 = np.triu_indices(4)

    >>> iu2 = np.triu_indices(4, 2)



    Here is how they can be used with a sample array:



    >>> a = np.arange(16).reshape(4, 4)

    >>> a

    array([[ 0,  1,  2,  3],

           [ 4,  5,  6,  7],

           [ 8,  9, 10, 11],

           [12, 13, 14, 15]])



    Both for indexing:



    >>> a[iu1]

    array([ 0,  1,  2, ..., 10, 11, 15])



    And for assigning values:



    >>> a[iu1] = -1

    >>> a

    array([[-1, -1, -1, -1],

           [ 4, -1, -1, -1],

           [ 8,  9, -1, -1],

           [12, 13, 14, -1]])



    These cover only a small part of the whole array (two diagonals right

    of the main one):



    >>> a[iu2] = -10

    >>> a

    array([[ -1,  -1, -10, -10],

           [  4,  -1,  -1, -10],

           [  8,   9,  -1,  -1],

           [ 12,  13,  14,  -1]])



    """
    tri_ = ~tri(n, m, k=k - 1, dtype=bool)

    return tuple(broadcast_to(inds, tri_.shape)[tri_]
                 for inds in indices(tri_.shape, sparse=True))


@array_function_dispatch(_trilu_indices_form_dispatcher)
def triu_indices_from(arr, k=0):
    """

    Return the indices for the upper-triangle of arr.



    See `triu_indices` for full details.



    Parameters

    ----------

    arr : ndarray, shape(N, N)

        The indices will be valid for square arrays.

    k : int, optional

        Diagonal offset (see `triu` for details).



    Returns

    -------

    triu_indices_from : tuple, shape(2) of ndarray, shape(N)

        Indices for the upper-triangle of `arr`.



    See Also

    --------

    triu_indices, triu



    Notes

    -----

    .. versionadded:: 1.4.0



    """
    if arr.ndim != 2:
        raise ValueError("input array must be 2-d")
    return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1])