Efficient way to apply function to each 2D slice of 3D numpy array

Question

I want to apply a function that takes a 2D array (and returns one of the same shape) to each 2D slice of a 3D array. What's an efficient way of doing this? numpy.fromiter returns a 1D array and numpy.fromfunction needs to be applied to each coordinate individually.

Currently I am doing

foo = np.array([func(arg, bar2D) for bar2D in bar3D])

This gives me what I want, but the list comprehension is very slow. Also, func is a 1D derivative with particular boundary conditions. numpy.gradient only seems to do N-D derivatives with N the dimension of the array, but maybe there is another routine that will do the whole thing for me?

Edit: The list comprehension works, but I'm looking for a faster way of doing it. bar3D can be large, up to (500,500,1000). All the numpy routines I've found for applying functions to arrays seem to assume either the function or the array are 1D.

Answer 1

I don't know of any generic way to apply functions to N-D slices of arrays. But there are two ways to get around it.

If what you want to do is apply a 1D derivative on each row or column of each 2D-slice, this is equivalent to applying the derivative to each 1D slice, and you can use np.apply_along_axis:

values = np.arange(4)*np.arange(3)[:, None]+np.arange(2)[:, None, None]*2
>>> array([[[0, 0, 0, 0],
            [0, 1, 2, 3],
            [0, 2, 4, 6]],

       [[2, 2, 2, 2],
        [2, 3, 4, 5],
        [2, 4, 6, 8]]])

np.apply_along_axis(np.gradient, 2, values)
>>> array([[[ 0.,  0.,  0.,  0.],
            [ 1.,  1.,  1.,  1.],
            [ 2.,  2.,  2.,  2.]],

           [[ 0.,  0.,  0.,  0.],
            [ 1.,  1.,  1.,  1.],
            [ 2.,  2.,  2.,  2.]]])

This differentiates the rows of each 2D slice. To differantiate each column do np.apply_along_axis(np.gradient, 2, values)

If you want to do something that requires two dimensions, you can usually get it through broadcasting and axis parameters. If for instance you want V[i, j] = sqrt((V[i,j]-V[i, j-1])^2+V[i, j]-V[i-1, j])^2 for each slice V you can do:

xdiffs = np.zeros_like(values) 
xdiffs[:, 1:, :]= np.diff(values, axis=1) 

ydiffs = np.zeros_like(values)
ydiffs[:, :, 1:] = np.diff(values, axis=2)

diffnorms = np.linalg.norm(xdiffs, ydiffs)

>>> array(
  [[[ 0.        ,  0.        ,  0.        ,  0.        ],
    [ 0.        ,  1.41421356,  2.23606798,  3.16227766],
    [ 0.        ,  2.23606798,  2.82842712,  3.60555128]],

   [[ 0.        ,  0.        ,  0.        ,  0.        ],
    [ 0.        ,  1.41421356,  2.23606798,  3.16227766],
    [ 0.        ,  2.23606798,  2.82842712,  3.60555128]]])

It's a bit cumbersome to get the dimensions right, but it will usually be the most efficient solution.

This examples uses zeros at the boundries, if you need something else, you need to set normdiff[:, :, 0] and normdiff[:, 0, :] to the correct boundry values.