comparison DEPENDENCIES/mingw32/Python27/Lib/site-packages/numpy/lib/index_tricks.py @ 87:2a2c65a20a8b

Add Python libs and headers

+from __future__ import division, absolute_import, print_function
+
+import sys
+import math
+
+import numpy.core.numeric as _nx
+from numpy.core.numeric import (
+    asarray, ScalarType, array, alltrue, cumprod, arange
+    )
+from numpy.core.numerictypes import find_common_type
+
+from . import function_base
+import numpy.matrixlib as matrix
+from .function_base import diff
+from numpy.lib._compiled_base import ravel_multi_index, unravel_index
+from numpy.lib.stride_tricks import as_strided
+
+makemat = matrix.matrix
+
+
+__all__ = [
+    'ravel_multi_index', 'unravel_index', 'mgrid', 'ogrid', 'r_', 'c_',
+    's_', 'index_exp', 'ix_', 'ndenumerate', 'ndindex', 'fill_diagonal',
+    'diag_indices', 'diag_indices_from'
+    ]
+
+
+def ix_(*args):
+    """
+    Construct an open mesh from multiple sequences.
+
+    This function takes N 1-D sequences and returns N outputs with N
+    dimensions each, such that the shape is 1 in all but one dimension
+    and the dimension with the non-unit shape value cycles through all
+    N dimensions.
+
+    Using `ix_` one can quickly construct index arrays that will index
+    the cross product. ``a[np.ix_([1,3],[2,5])]`` returns the array
+    ``[[a[1,2] a[1,5]], [a[3,2] a[3,5]]]``.
+
+    Parameters
+    ----------
+    args : 1-D sequences
+
+    Returns
+    -------
+    out : tuple of ndarrays
+        N arrays with N dimensions each, with N the number of input
+        sequences. Together these arrays form an open mesh.
+
+    See Also
+    --------
+    ogrid, mgrid, meshgrid
+
+    Examples
+    --------
+    >>> a = np.arange(10).reshape(2, 5)
+    >>> a
+    array([[0, 1, 2, 3, 4],
+           [5, 6, 7, 8, 9]])
+    >>> ixgrid = np.ix_([0,1], [2,4])
+    >>> ixgrid
+    (array([[0],
+           [1]]), array([[2, 4]]))
+    >>> ixgrid[0].shape, ixgrid[1].shape
+    ((2, 1), (1, 2))
+    >>> a[ixgrid]
+    array([[2, 4],
+           [7, 9]])
+
+    """
+    out = []
+    nd = len(args)
+    baseshape = [1]*nd
+    for k in range(nd):
+        new = _nx.asarray(args[k])
+        if (new.ndim != 1):
+            raise ValueError("Cross index must be 1 dimensional")
+        if issubclass(new.dtype.type, _nx.bool_):
+            new = new.nonzero()[0]
+        baseshape[k] = len(new)
+        new = new.reshape(tuple(baseshape))
+        out.append(new)
+        baseshape[k] = 1
+    return tuple(out)
+
+class nd_grid(object):
+    """
+    Construct a multi-dimensional "meshgrid".
+
+    ``grid = nd_grid()`` creates an instance which will return a mesh-grid
+    when indexed.  The dimension and number of the output arrays are equal
+    to the number of indexing dimensions.  If the step length is not a
+    complex number, then the stop is not inclusive.
+
+    However, if the step length is a **complex number** (e.g. 5j), then the
+    integer part of its magnitude is interpreted as specifying the
+    number of points to create between the start and stop values, where
+    the stop value **is inclusive**.
+
+    If instantiated with an argument of ``sparse=True``, the mesh-grid is
+    open (or not fleshed out) so that only one-dimension of each returned
+    argument is greater than 1.
+
+    Parameters
+    ----------
+    sparse : bool, optional
+        Whether the grid is sparse or not. Default is False.
+
+    Notes
+    -----
+    Two instances of `nd_grid` are made available in the NumPy namespace,
+    `mgrid` and `ogrid`::
+
+        mgrid = nd_grid(sparse=False)
+        ogrid = nd_grid(sparse=True)
+
+    Users should use these pre-defined instances instead of using `nd_grid`
+    directly.
+
+    Examples
+    --------
+    >>> mgrid = np.lib.index_tricks.nd_grid()
+    >>> mgrid[0:5,0:5]
+    array([[[0, 0, 0, 0, 0],
+            [1, 1, 1, 1, 1],
+            [2, 2, 2, 2, 2],
+            [3, 3, 3, 3, 3],
+            [4, 4, 4, 4, 4]],
+           [[0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4]]])
+    >>> mgrid[-1:1:5j]
+    array([-1. , -0.5,  0. ,  0.5,  1. ])
+
+    >>> ogrid = np.lib.index_tricks.nd_grid(sparse=True)
+    >>> ogrid[0:5,0:5]
+    [array([[0],
+            [1],
+            [2],
+            [3],
+            [4]]), array([[0, 1, 2, 3, 4]])]
+
+    """
+
+    def __init__(self, sparse=False):
+        self.sparse = sparse
+
+    def __getitem__(self, key):
+        try:
+            size = []
+            typ = int
+            for k in range(len(key)):
+                step = key[k].step
+                start = key[k].start
+                if start is None:
+                    start = 0
+                if step is None:
+                    step = 1
+                if isinstance(step, complex):
+                    size.append(int(abs(step)))
+                    typ = float
+                else:
+                    size.append(
+                        int(math.ceil((key[k].stop - start)/(step*1.0))))
+                if (isinstance(step, float) or
+                        isinstance(start, float) or
+                        isinstance(key[k].stop, float)):
+                    typ = float
+            if self.sparse:
+                nn = [_nx.arange(_x, dtype=_t)
+                        for _x, _t in zip(size, (typ,)*len(size))]
+            else:
+                nn = _nx.indices(size, typ)
+            for k in range(len(size)):
+                step = key[k].step
+                start = key[k].start
+                if start is None:
+                    start = 0
+                if step is None:
+                    step = 1
+                if isinstance(step, complex):
+                    step = int(abs(step))
+                    if step != 1:
+                        step = (key[k].stop - start)/float(step-1)
+                nn[k] = (nn[k]*step+start)
+            if self.sparse:
+                slobj = [_nx.newaxis]*len(size)
+                for k in range(len(size)):
+                    slobj[k] = slice(None, None)
+                    nn[k] = nn[k][slobj]
+                    slobj[k] = _nx.newaxis
+            return nn
+        except (IndexError, TypeError):
+            step = key.step
+            stop = key.stop
+            start = key.start
+            if start is None:
+                start = 0
+            if isinstance(step, complex):
+                step = abs(step)
+                length = int(step)
+                if step != 1:
+                    step = (key.stop-start)/float(step-1)
+                stop = key.stop + step
+                return _nx.arange(0, length, 1, float)*step + start
+            else:
+                return _nx.arange(start, stop, step)
+
+    def __getslice__(self, i, j):
+        return _nx.arange(i, j)
+
+    def __len__(self):
+        return 0
+
+mgrid = nd_grid(sparse=False)
+ogrid = nd_grid(sparse=True)
+mgrid.__doc__ = None  # set in numpy.add_newdocs
+ogrid.__doc__ = None  # set in numpy.add_newdocs
+
+class AxisConcatenator(object):
+    """
+    Translates slice objects to concatenation along an axis.
+
+    For detailed documentation on usage, see `r_`.
+
+    """
+
+    def _retval(self, res):
+        if self.matrix:
+            oldndim = res.ndim
+            res = makemat(res)
+            if oldndim == 1 and self.col:
+                res = res.T
+        self.axis = self._axis
+        self.matrix = self._matrix
+        self.col = 0
+        return res
+
+    def __init__(self, axis=0, matrix=False, ndmin=1, trans1d=-1):
+        self._axis = axis
+        self._matrix = matrix
+        self.axis = axis
+        self.matrix = matrix
+        self.col = 0
+        self.trans1d = trans1d
+        self.ndmin = ndmin
+
+    def __getitem__(self, key):
+        trans1d = self.trans1d
+        ndmin = self.ndmin
+        if isinstance(key, str):
+            frame = sys._getframe().f_back
+            mymat = matrix.bmat(key, frame.f_globals, frame.f_locals)
+            return mymat
+        if not isinstance(key, tuple):
+            key = (key,)
+        objs = []
+        scalars = []
+        arraytypes = []
+        scalartypes = []
+        for k in range(len(key)):
+            scalar = False
+            if isinstance(key[k], slice):
+                step = key[k].step
+                start = key[k].start
+                stop = key[k].stop
+                if start is None:
+                    start = 0
+                if step is None:
+                    step = 1
+                if isinstance(step, complex):
+                    size = int(abs(step))
+                    newobj = function_base.linspace(start, stop, num=size)
+                else:
+                    newobj = _nx.arange(start, stop, step)
+                if ndmin > 1:
+                    newobj = array(newobj, copy=False, ndmin=ndmin)
+                    if trans1d != -1:
+                        newobj = newobj.swapaxes(-1, trans1d)
+            elif isinstance(key[k], str):
+                if k != 0:
+                    raise ValueError("special directives must be the "
+                            "first entry.")
+                key0 = key[0]
+                if key0 in 'rc':
+                    self.matrix = True
+                    self.col = (key0 == 'c')
+                    continue
+                if ',' in key0:
+                    vec = key0.split(',')
+                    try:
+                        self.axis, ndmin = \
+                                   [int(x) for x in vec[:2]]
+                        if len(vec) == 3:
+                            trans1d = int(vec[2])
+                        continue
+                    except:
+                        raise ValueError("unknown special directive")
+                try:
+                    self.axis = int(key[k])
+                    continue
+                except (ValueError, TypeError):
+                    raise ValueError("unknown special directive")
+            elif type(key[k]) in ScalarType:
+                newobj = array(key[k], ndmin=ndmin)
+                scalars.append(k)
+                scalar = True
+                scalartypes.append(newobj.dtype)
+            else:
+                newobj = key[k]
+                if ndmin > 1:
+                    tempobj = array(newobj, copy=False, subok=True)
+                    newobj = array(newobj, copy=False, subok=True,
+                                   ndmin=ndmin)
+                    if trans1d != -1 and tempobj.ndim < ndmin:
+                        k2 = ndmin-tempobj.ndim
+                        if (trans1d < 0):
+                            trans1d += k2 + 1
+                        defaxes = list(range(ndmin))
+                        k1 = trans1d
+                        axes = defaxes[:k1] + defaxes[k2:] + \
+                               defaxes[k1:k2]
+                        newobj = newobj.transpose(axes)
+                    del tempobj
+            objs.append(newobj)
+            if not scalar and isinstance(newobj, _nx.ndarray):
+                arraytypes.append(newobj.dtype)
+
+        #  Esure that scalars won't up-cast unless warranted
+        final_dtype = find_common_type(arraytypes, scalartypes)
+        if final_dtype is not None:
+            for k in scalars:
+                objs[k] = objs[k].astype(final_dtype)
+
+        res = _nx.concatenate(tuple(objs), axis=self.axis)
+        return self._retval(res)
+
+    def __getslice__(self, i, j):
+        res = _nx.arange(i, j)
+        return self._retval(res)
+
+    def __len__(self):
+        return 0
+
+# separate classes are used here instead of just making r_ = concatentor(0),
+# etc. because otherwise we couldn't get the doc string to come out right
+# in help(r_)
+
+class RClass(AxisConcatenator):
+    """
+    Translates slice objects to concatenation along the first axis.
+
+    This is a simple way to build up arrays quickly. There are two use cases.
+
+    1. If the index expression contains comma separated arrays, then stack
+       them along their first axis.
+    2. If the index expression contains slice notation or scalars then create
+       a 1-D array with a range indicated by the slice notation.
+
+    If slice notation is used, the syntax ``start:stop:step`` is equivalent
+    to ``np.arange(start, stop, step)`` inside of the brackets. However, if
+    ``step`` is an imaginary number (i.e. 100j) then its integer portion is
+    interpreted as a number-of-points desired and the start and stop are
+    inclusive. In other words ``start:stop:stepj`` is interpreted as
+    ``np.linspace(start, stop, step, endpoint=1)`` inside of the brackets.
+    After expansion of slice notation, all comma separated sequences are
+    concatenated together.
+
+    Optional character strings placed as the first element of the index
+    expression can be used to change the output. The strings 'r' or 'c' result
+    in matrix output. If the result is 1-D and 'r' is specified a 1 x N (row)
+    matrix is produced. If the result is 1-D and 'c' is specified, then a N x 1
+    (column) matrix is produced. If the result is 2-D then both provide the
+    same matrix result.
+
+    A string integer specifies which axis to stack multiple comma separated
+    arrays along. A string of two comma-separated integers allows indication
+    of the minimum number of dimensions to force each entry into as the
+    second integer (the axis to concatenate along is still the first integer).
+
+    A string with three comma-separated integers allows specification of the
+    axis to concatenate along, the minimum number of dimensions to force the
+    entries to, and which axis should contain the start of the arrays which
+    are less than the specified number of dimensions. In other words the third
+    integer allows you to specify where the 1's should be placed in the shape
+    of the arrays that have their shapes upgraded. By default, they are placed
+    in the front of the shape tuple. The third argument allows you to specify
+    where the start of the array should be instead. Thus, a third argument of
+    '0' would place the 1's at the end of the array shape. Negative integers
+    specify where in the new shape tuple the last dimension of upgraded arrays
+    should be placed, so the default is '-1'.
+
+    Parameters
+    ----------
+    Not a function, so takes no parameters
+
+
+    Returns
+    -------
+    A concatenated ndarray or matrix.
+
+    See Also
+    --------
+    concatenate : Join a sequence of arrays together.
+    c_ : Translates slice objects to concatenation along the second axis.
+
+    Examples
+    --------
+    >>> np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])]
+    array([1, 2, 3, 0, 0, 4, 5, 6])
+    >>> np.r_[-1:1:6j, [0]*3, 5, 6]
+    array([-1. , -0.6, -0.2,  0.2,  0.6,  1. ,  0. ,  0. ,  0. ,  5. ,  6. ])
+
+    String integers specify the axis to concatenate along or the minimum
+    number of dimensions to force entries into.
+
+    >>> a = np.array([[0, 1, 2], [3, 4, 5]])
+    >>> np.r_['-1', a, a] # concatenate along last axis
+    array([[0, 1, 2, 0, 1, 2],
+           [3, 4, 5, 3, 4, 5]])
+    >>> np.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, dim>=2
+    array([[1, 2, 3],
+           [4, 5, 6]])
+
+    >>> np.r_['0,2,0', [1,2,3], [4,5,6]]
+    array([[1],
+           [2],
+           [3],
+           [4],
+           [5],
+           [6]])
+    >>> np.r_['1,2,0', [1,2,3], [4,5,6]]
+    array([[1, 4],
+           [2, 5],
+           [3, 6]])
+
+    Using 'r' or 'c' as a first string argument creates a matrix.
+
+    >>> np.r_['r',[1,2,3], [4,5,6]]
+    matrix([[1, 2, 3, 4, 5, 6]])
+
+    """
+
+    def __init__(self):
+        AxisConcatenator.__init__(self, 0)
+
+r_ = RClass()
+
+class CClass(AxisConcatenator):
+    """
+    Translates slice objects to concatenation along the second axis.
+
+    This is short-hand for ``np.r_['-1,2,0', index expression]``, which is
+    useful because of its common occurrence. In particular, arrays will be
+    stacked along their last axis after being upgraded to at least 2-D with
+    1's post-pended to the shape (column vectors made out of 1-D arrays).
+
+    For detailed documentation, see `r_`.
+
+    Examples
+    --------
+    >>> np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])]
+    array([[1, 2, 3, 0, 0, 4, 5, 6]])
+
+    """
+
+    def __init__(self):
+        AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0)
+
+c_ = CClass()
+
+class ndenumerate(object):
+    """
+    Multidimensional index iterator.
+
+    Return an iterator yielding pairs of array coordinates and values.
+
+    Parameters
+    ----------
+    a : ndarray
+      Input array.
+
+    See Also
+    --------
+    ndindex, flatiter
+
+    Examples
+    --------
+    >>> a = np.array([[1, 2], [3, 4]])
+    >>> for index, x in np.ndenumerate(a):
+    ...     print index, x
+    (0, 0) 1
+    (0, 1) 2
+    (1, 0) 3
+    (1, 1) 4
+
+    """
+
+    def __init__(self, arr):
+        self.iter = asarray(arr).flat
+
+    def __next__(self):
+        """
+        Standard iterator method, returns the index tuple and array value.
+
+        Returns
+        -------
+        coords : tuple of ints
+            The indices of the current iteration.
+        val : scalar
+            The array element of the current iteration.
+
+        """
+        return self.iter.coords, next(self.iter)
+
+    def __iter__(self):
+        return self
+
+    next = __next__
+
+
+class ndindex(object):
+    """
+    An N-dimensional iterator object to index arrays.
+
+    Given the shape of an array, an `ndindex` instance iterates over
+    the N-dimensional index of the array. At each iteration a tuple
+    of indices is returned, the last dimension is iterated over first.
+
+    Parameters
+    ----------
+    `*args` : ints
+      The size of each dimension of the array.
+
+    See Also
+    --------
+    ndenumerate, flatiter
+
+    Examples
+    --------
+    >>> for index in np.ndindex(3, 2, 1):
+    ...     print index
+    (0, 0, 0)
+    (0, 1, 0)
+    (1, 0, 0)
+    (1, 1, 0)
+    (2, 0, 0)
+    (2, 1, 0)
+
+    """
+
+    def __init__(self, *shape):
+        if len(shape) == 1 and isinstance(shape[0], tuple):
+            shape = shape[0]
+        x = as_strided(_nx.zeros(1), shape=shape,
+                       strides=_nx.zeros_like(shape))
+        self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'],
+                              order='C')
+
+    def __iter__(self):
+        return self
+
+    def ndincr(self):
+        """
+        Increment the multi-dimensional index by one.
+
+        This method is for backward compatibility only: do not use.
+        """
+        next(self)
+
+    def __next__(self):
+        """
+        Standard iterator method, updates the index and returns the index
+        tuple.
+
+        Returns
+        -------
+        val : tuple of ints
+            Returns a tuple containing the indices of the current
+            iteration.
+
+        """
+        next(self._it)
+        return self._it.multi_index
+
+    next = __next__
+
+
+# You can do all this with slice() plus a few special objects,
+# but there's a lot to remember. This version is simpler because
+# it uses the standard array indexing syntax.
+#
+# Written by Konrad Hinsen <[e-mail]>
+# last revision: 1999-7-23
+#
+# Cosmetic changes by T. Oliphant 2001
+#
+#
+
+class IndexExpression(object):
+    """
+    A nicer way to build up index tuples for arrays.
+
+    .. note::
+       Use one of the two predefined instances `index_exp` or `s_`
+       rather than directly using `IndexExpression`.
+
+    For any index combination, including slicing and axis insertion,
+    ``a[indices]`` is the same as ``a[np.index_exp[indices]]`` for any
+    array `a`. However, ``np.index_exp[indices]`` can be used anywhere
+    in Python code and returns a tuple of slice objects that can be
+    used in the construction of complex index expressions.
+
+    Parameters
+    ----------
+    maketuple : bool
+        If True, always returns a tuple.
+
+    See Also
+    --------
+    index_exp : Predefined instance that always returns a tuple:
+       `index_exp = IndexExpression(maketuple=True)`.
+    s_ : Predefined instance without tuple conversion:
+       `s_ = IndexExpression(maketuple=False)`.
+
+    Notes
+    -----
+    You can do all this with `slice()` plus a few special objects,
+    but there's a lot to remember and this version is simpler because
+    it uses the standard array indexing syntax.
+
+    Examples
+    --------
+    >>> np.s_[2::2]
+    slice(2, None, 2)
+    >>> np.index_exp[2::2]
+    (slice(2, None, 2),)
+
+    >>> np.array([0, 1, 2, 3, 4])[np.s_[2::2]]
+    array([2, 4])
+
+    """
+
+    def __init__(self, maketuple):
+        self.maketuple = maketuple
+
+    def __getitem__(self, item):
+        if self.maketuple and not isinstance(item, tuple):
+            return (item,)
+        else:
+            return item
+
+index_exp = IndexExpression(maketuple=True)
+s_ = IndexExpression(maketuple=False)
+
+# End contribution from Konrad.
+
+
+# The following functions complement those in twodim_base, but are
+# applicable to N-dimensions.
+
+def fill_diagonal(a, val, wrap=False):
+    """Fill the main diagonal of the given array of any dimensionality.
+
+    For an array `a` with ``a.ndim > 2``, the diagonal is the list of
+    locations with indices ``a[i, i, ..., i]`` all identical. This function
+    modifies the input array in-place, it does not return a value.
+
+    Parameters
+    ----------
+    a : array, at least 2-D.
+      Array whose diagonal is to be filled, it gets modified in-place.
+
+    val : scalar
+      Value to be written on the diagonal, its type must be compatible with
+      that of the array a.
+
+    wrap : bool
+      For tall matrices in NumPy version up to 1.6.2, the
+      diagonal "wrapped" after N columns. You can have this behavior
+      with this option. This affect only tall matrices.
+
+    See also
+    --------
+    diag_indices, diag_indices_from
+
+    Notes
+    -----
+    .. versionadded:: 1.4.0
+
+    This functionality can be obtained via `diag_indices`, but internally
+    this version uses a much faster implementation that never constructs the
+    indices and uses simple slicing.
+
+    Examples
+    --------
+    >>> a = np.zeros((3, 3), int)
+    >>> np.fill_diagonal(a, 5)
+    >>> a
+    array([[5, 0, 0],
+           [0, 5, 0],
+           [0, 0, 5]])
+
+    The same function can operate on a 4-D array:
+
+    >>> a = np.zeros((3, 3, 3, 3), int)
+    >>> np.fill_diagonal(a, 4)
+
+    We only show a few blocks for clarity:
+
+    >>> a[0, 0]
+    array([[4, 0, 0],
+           [0, 0, 0],
+           [0, 0, 0]])
+    >>> a[1, 1]
+    array([[0, 0, 0],
+           [0, 4, 0],
+           [0, 0, 0]])
+    >>> a[2, 2]
+    array([[0, 0, 0],
+           [0, 0, 0],
+           [0, 0, 4]])
+
+    # tall matrices no wrap
+    >>> a = np.zeros((5, 3),int)
+    >>> fill_diagonal(a, 4)
+    array([[4, 0, 0],
+           [0, 4, 0],
+           [0, 0, 4],
+           [0, 0, 0],
+           [0, 0, 0]])
+
+    # tall matrices wrap
+    >>> a = np.zeros((5, 3),int)
+    >>> fill_diagonal(a, 4)
+    array([[4, 0, 0],
+           [0, 4, 0],
+           [0, 0, 4],
+           [0, 0, 0],
+           [4, 0, 0]])
+
+    # wide matrices
+    >>> a = np.zeros((3, 5),int)
+    >>> fill_diagonal(a, 4)
+    array([[4, 0, 0, 0, 0],
+           [0, 4, 0, 0, 0],
+           [0, 0, 4, 0, 0]])
+
+    """
+    if a.ndim < 2:
+        raise ValueError("array must be at least 2-d")
+    end = None
+    if a.ndim == 2:
+        # Explicit, fast formula for the common case.  For 2-d arrays, we
+        # accept rectangular ones.
+        step = a.shape[1] + 1
+        #This is needed to don't have tall matrix have the diagonal wrap.
+        if not wrap:
+            end = a.shape[1] * a.shape[1]
+    else:
+        # For more than d=2, the strided formula is only valid for arrays with
+        # all dimensions equal, so we check first.
+        if not alltrue(diff(a.shape) == 0):
+            raise ValueError("All dimensions of input must be of equal length")
+        step = 1 + (cumprod(a.shape[:-1])).sum()
+
+    # Write the value out into the diagonal.
+    a.flat[:end:step] = val
+
+
+def diag_indices(n, ndim=2):
+    """
+    Return the indices to access the main diagonal of an array.
+
+    This returns a tuple of indices that can be used to access the main
+    diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape
+    (n, n, ..., n). For ``a.ndim = 2`` this is the usual diagonal, for
+    ``a.ndim > 2`` this is the set of indices to access ``a[i, i, ..., i]``
+    for ``i = [0..n-1]``.
+
+    Parameters
+    ----------
+    n : int
+      The size, along each dimension, of the arrays for which the returned
+      indices can be used.
+
+    ndim : int, optional
+      The number of dimensions.
+
+    See also
+    --------
+    diag_indices_from
+
+    Notes
+    -----
+    .. versionadded:: 1.4.0
+
+    Examples
+    --------
+    Create a set of indices to access the diagonal of a (4, 4) array:
+
+    >>> di = np.diag_indices(4)
+    >>> di
+    (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
+    >>> 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]])
+    >>> a[di] = 100
+    >>> a
+    array([[100,   1,   2,   3],
+           [  4, 100,   6,   7],
+           [  8,   9, 100,  11],
+           [ 12,  13,  14, 100]])
+
+    Now, we create indices to manipulate a 3-D array:
+
+    >>> d3 = np.diag_indices(2, 3)
+    >>> d3
+    (array([0, 1]), array([0, 1]), array([0, 1]))
+
+    And use it to set the diagonal of an array of zeros to 1:
+
+    >>> a = np.zeros((2, 2, 2), dtype=np.int)
+    >>> a[d3] = 1
+    >>> a
+    array([[[1, 0],
+            [0, 0]],
+           [[0, 0],
+            [0, 1]]])
+
+    """
+    idx = arange(n)
+    return (idx,) * ndim
+
+
+def diag_indices_from(arr):
+    """
+    Return the indices to access the main diagonal of an n-dimensional array.
+
+    See `diag_indices` for full details.
+
+    Parameters
+    ----------
+    arr : array, at least 2-D
+
+    See Also
+    --------
+    diag_indices
+
+    Notes
+    -----
+    .. versionadded:: 1.4.0
+
+    """
+
+    if not arr.ndim >= 2:
+        raise ValueError("input array must be at least 2-d")
+    # For more than d=2, the strided formula is only valid for arrays with
+    # all dimensions equal, so we check first.
+    if not alltrue(diff(arr.shape) == 0):
+        raise ValueError("All dimensions of input must be of equal length")
+
+    return diag_indices(arr.shape[0], arr.ndim)