Numpy Vectorized Function Over Successive 2d Slices

Question

I have a 3D numpy array. I would like to form a new 3d array by executing a function on successive 2d slices along an axis, and stacking the resulting slices together. Clearly there are many ways to do this; I'd like to do it in the most concise way possible. I'd think this would be possible with numpy.vectorize, but this seems to produce a function that iterates over every value in my array, rather than 2D slices taken by moving along the first axis.

Basically, I want code that looks something like this:

new3dmat = np.vectorize(func2dmat)(my3dmat)

And accomplishes the same thing as this:

new3dmat = np.empty_like(my3dmat)
for i in range(my3dmat.shape[0]):
  new3dmat[i] = func2dmat(my3dmat[i])

How can I accomplish this?

Answer 1

I am afraid something like below is the best compromise between conciseness and performance. apply_along_axis does not take multiple axes, unfortunately.

new3dmat = np.array([func2dmat(slice) for slice in my3dmat])

It isn't ideal in terms of extra allocations and so on, but unless .shape[0] is big relative to .size, the extra overhead should be minimal.