Iterate over numpy with index (numpy equivalent of python enumerate)

Question

I\'m trying to create a function that will calculate the lattice distance (number of horizontal and vertical steps) between elements in a multi-dimensional numpy array. For

Answer 1

You can do it using np.ndenumerate but generally you don't need to iterate over an array.

You can simply create a meshgrid (or open grid) to get all indices at once and you can then process them (vectorized) much faster.

For example

>>> x, y = np.mgrid[slice(A.shape[0]), slice(A.shape[1])]
>>> x
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2]])
>>> y
array([[0, 1, 2],
       [0, 1, 2],
       [0, 1, 2]])

and these can be processed like any other array. So if your function that needs the indices can be vectorized you shouldn't do the manual loop!

For example to calculate the lattice distance for each point to a point say (2, 3):

>>> abs(x - 2) + abs(y - 3)
array([[5, 4, 3],
       [4, 3, 2],
       [3, 2, 1]])

For distances an ogrid would be faster. Just replace np.mgrid with np.ogrid:

>>> x, y = np.ogrid[slice(A.shape[0]), slice(A.shape[1])]
>>> np.hypot(x - 2, y - 3)  # cartesian distance this time! :-)
array([[ 3.60555128,  2.82842712,  2.23606798],
       [ 3.16227766,  2.23606798,  1.41421356],
       [ 3.        ,  2.        ,  1.        ]])