Python apply_along_axis of multiple arrays

Question

If I have a function, f(x) which takes a single 1d-array as argument and produces a 1d-array as output, I can use numpy.apply_along_axis to apply the function, to each row of a 2d-array X whose rows are valid arguments for f.

Now I want to do the analogous thing with a function that takes two arguments. E.g. I have a function f(x, y) which takes tow 1d-arrays as arguments and I also have two 2d-arrays X, Y both with n rows. I want to apply f to each pair of rows, producing an array which has again n rows.

How to accomplish this in an efficient way?

I am also interested in variants, where f takes more arguments or higher dimensional arrays are involved:

For example f could take 3 arrays x, y, z of shape (2,2); (3,); (5,) and produce a result of shape (4,4).

I have X, Y, Z of shapes (50, 100, 2, 2); (50, 100, 3); (50, 100, 5) and want a result of shape (50, 100, 4, 4)

Answer 1

Looking at the code for numpy.apply_along_axis I see that it just iterates over the other dimensions, applying your function to each 'row'. There's extra code that allows for dimensions about 2. But for 2d X it boils down to:

result = np.empty_like(X)
for i, x in enumerate(X):
    result[i] = func1d(x)

There's also code to deduce what shape the result should have. For example if func1d is np.sum, then result will be 1d, not 2d like the input.

So there's not special 'efficiency' in this function. An extension to multiple inputs could be a normal Python zip:

 result = np.empty_like(X)
 for i,(x,y) in enumerate(zip(X,Y)):
     result[i] = func1d(x,y)

---

np.ndindex is handy tool for generating indices. It's worth looking at its code. It uses the general purpose numpy iterator, np.nditer, which see: http://docs.scipy.org/doc/numpy/reference/arrays.nditer.html

For example f could take 3 arrays x, y, z of shape (2,2); (3,); (5,) and produce a result of shape (4,4).

I have X, Y, Z of shapes (50, 100, 2, 2); (50, 100, 3); (50, 100, 5) and want a result of shape (50, 100, 4, 4)

for i,j in np.ndindex(50,100):
    result[i,j,:,:] = f(X[i,j,:,:], Y[i,j,:,:], Z[i,j,:,:])

The ':' aren't necessary, but make it clear that we are indexing on 2 of the dimensions, and slicing the rest. They would be required if you wanted to iterate on the 1 and 3rd dimensions, and slice the 2nd.