Best way to apply a function to a slice of a 3d numpy array

Question

Suppose I have something like myArray.shape == (100, 80, 2)

I want to do something like this: numpy.apply_along_axis(function, 0, myArray) where function uses both items on the axis=2 axis of myArray, but I know numpy.apply_along_axis only works for 1D slices.

My question is: Is there a generic way to go about acting a function to 2D slices without having to use a loop or does it depend on how I have function defined? And if so, what would be the most efficient way of doing this?

Is it possible to use numpy.apply_along_axis to act on one 1D slice and zip each element in the other slice to each element in the first slice somehow? Would it help to restructure myArray?

Note: This question did not answer my question, so please don't mark as duplicate.

"does it depend on how I have function defined?"... Yes, very much. — cs95, Mar 25 '18 at 19:43
`apply_along` uses loops. It's nkt entirely clear what you trying to do. Show us with an explicit loop. — hpaulj, Mar 25 '18 at 20:02

score 7 · Accepted Answer · answered Mar 25 '18 at 20:52

Define a simple function that takes a 2d array, and returns a scalar

In [54]: def foo(x):
    ...:     assert(x.ndim == 2)
    ...:     return x.mean()
    ...: 
In [55]: X = np.arange(24).reshape(2,3,4)

It's not entirely clear how you want to iterate on the 3d array, but let's assume it's on the first axis. The straight forward list comprehension approach is:

In [56]: [foo(x) for x in X]
Out[56]: [5.5, 17.5]

vectorize normally feeds scalars to the function, but the newer versions have signature parameter that allow us to use it as:

In [58]: f = np.vectorize(foo, signature='(n,m)->()')
In [59]: f(X)
Out[59]: array([ 5.5, 17.5])

The original vectorize does not promise any speed up, and the signature version is even a bit slower.

apply_along_axis just hides iteration. Even though it only operates on 1d arrays, we can use it with a bit of reshaping:

In [62]: np.apply_along_axis(lambda x: foo(x.reshape(3,4)), 1, X.reshape(2,-1))
Out[62]: array([ 5.5, 17.5])

As long as you are only iterating on one axis, the list comprehension approach is both fastest and easiest.

Best way to apply a function to a slice of a 3d numpy array

1 Answers1