векторизовать функцию, работающую на подмассиве ndarray

Question

У меня есть функция, действующая на каждый 2D-фрагмент 3D-массива. Как векторизовать функцию, чтобы избежать цикла для улучшения производительности? Например:

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

Это не может быть векторизованы без реализовав interp_2d. Однако, если предположить, что interp_2d - это некоторый тип интерполяции, то операция, вероятно, является линейной. То есть lambda z0: interp_2d(x0, y0, z0, x1, y1), вероятно, эквивалентен np.dot(M, z0), где M - некоторая (вероятно, редкая) матрица, которая зависит от x0, y0, x1 и y1. Прямо сейчас, вызывая функцию interp_2d, вы неявно пересчитываете эту матрицу при каждом вызове, даже если она одинакова каждый раз. Эффективнее выяснить, что такое матрица, и повторно применить ее к новому z0 много раз.

Вот очень простой пример 1D интерполяции:

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] 

Если n очень большой, ожидать скорость до хорошо над коэффициентом 100.