Vectorize a function acting on a ndarray subarray

Question

I have a function that acts on every 2D fragment of a 3D array. How to vectorize a function to avoid a loop to improve performance? For instance:

def interp_2d(x0,y0,z0,x1,y1): # x0, y0 and z0 are 2D array # x1 and y1 are 2D array # peform 2D interpolation return z1 # now I want to call the interp_2d for each 2D slice of z0_3d as following: for k in range(z0_3d.shape[2]): z1_3d[:,:,k]=interp_2d(x0, y0, z0_3d[:,:,k], x1, y1) 

Answer 1

This cannot be vectorized without overriding interp_2d . However, assuming interp_2d is some type of interpolation, the operation is probably linear. This is lambda z0: interp_2d(x0, y0, z0, x1, y1) , probably equivalent to np.dot(M, z0) , where M is some (probably sparse) matrix that depends on x0 , y0 , x1 and y1 . Right now, by calling the interp_2d function, you implicitly recalculate this matrix with every call, although the same thing every time. It is more efficient to find out what this matrix is ​​once and reapply it to the new z0 many times.

Here is a really trivial example of 1D interpolation:

 x0 = [0., 1.] x1 = 0.3 z0_2d = "some very long array with shape=(2, n)" def interp_1d(x0, z0, x1): """x0 and z0 are length 2, 1D arrays, x1 is a float between x0[0] and x0[1].""" delta_x = x0[1] - x0[0] w0 = (x1 - x0[0]) / delta_x w1 = (x0[1] - x1) / delta_x return w0 * z0[0] + w1 * z0[1] # The slow way. for i in range(n): z1_2d[i] = interp_1d(x0, z0_2d[:,i], x1) # Notice that the intermediate products w1 and w2 are the same on each # iteration but we recalculate them anyway. # The fast way. def interp_1d_weights(x0, x1): delta_x = x0[1] - x0[0] w0 = (x1 - x0[0]) / delta_x w1 = (x0[1] - x1) / delta_x return w0, w1 w0, w1 = interp_1d_weights(x0, x1) z1_2d = w0 * z0_2d[0,:] + w1 * z0_2d[1:0] 

If n very large, expect speed to be well above 100.