How to apply an index-dependent function to a numpy ndarray?

Question

So numpy ndarrays are quite handy in that you can just type in f(A) for any one-dimensional function f and any ndarray A and it will apply f element-wise. It is also, I was told, a very efficient way of applying a function to a ndarray and avoiding for loops. Avoid for loops, is what I have been told.

Turns out that now I need to apply a function f(A) that is not just one dimensional, but requires knowledge of the index tuple of each element in order to return the correct value for each element. Is there a way to avoid using for loops or explicit recursion and keep working with the direct function application on ndarrays under these circumstances? Or am I out of options?

Answer 1

Use numpy.meshgrid to generate coordinate matrices corresponding to index tuples of each element in the array. Then let your function also depend on these coordinates.

For example a is a three dimensional array.

x, y, z = np.meshgrid(np.arange(a.shape[0]), np.arange(a.shape[1]), np.arange(a.shape[2]), indexing='ij')

gives three arrays x, y, z which contains x, y and z coordinates at each location. The function on array a would then be extended by also giving the index arrays.

f(a, x, y, z)

Be careful with the order of the indices/directions. Check the options of indexing.