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