vectorize a function operated on subarray of a ndarray

Question

I have a function acted on each 2D slices of a 3D array. How to vectorize the function to avoid loop to improve the performace? For example:

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 can not be vectorized without reimplementing interp_2d. However, assuming interp_2d is some type of interpolation, then the operation is probably linear. That is lambda z0: interp_2d(x0, y0, z0, x1, y1) is 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 are implicitly recalculating this matrix on every call even though it is the same each time. It's more efficient to figure out what that matrix is once and reapply it to the new z0 many times.

Here is a really trivial 1D interpolation example:

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 is very large, expect a speed up of well over a factor of 100.