Быстрая интерполяция по 3D-массиву для 3D-начала x

Question

Эта проблема аналогична предыдущей заданной проблеме Fast interpolation over 3D array, но не может решить мою проблему.

У меня есть 4D-массив с размерами (время, высота, широта, долгота), обозначенное как y.shape=(nt, nalt, nlat, nlon). X - высота и изменение с (время, широта, долгота), что означает x.shape = (nt, nalt, nlat, nlon). Я хочу интерполировать на высоте для каждого (nt, nlat, nlon). Интерполированное x_new должно быть 1d, а не изменять (время, широта, долгота).

Я использую numpy.interp, то же, что и scipy.interpolate.interp1d, и подумайте об ответах на прежнем посту. Я не могу уменьшить петли с этими ответами.

я могу только так:

# y is a 4D ndarray 
# x is a 4D ndarray 
# new_y is a 4D array 
for i in range(nlon): 
    for j in range(nlat): 
     for k in range(nt): 
      y_new[k,:,j,i] = np.interp(new_x, x[k,:,j,i], y[k,:,j,i]) 

Эти петли делают это интерполяция слишком медленно расчета. У кого-нибудь есть хорошие идеи? Помощь будет высоко оценена.

Answer 1

Вот мое решение, используя numba, это примерно в 3 раза быстрее.

создать тестовые данные во-первых, x необходимость в порядке возрастания:

import numpy as np 
rows = 200000 
cols = 66 
new_cols = 69 
x = np.random.rand(rows, cols) 
x.sort(axis=-1) 
y = np.random.rand(rows, cols) 
nx = np.random.rand(new_cols) 
nx.sort() 

сделать 200000 раз интерп в NumPy:

%%time 
ny = np.empty((x.shape[0], len(nx))) 
for i in range(len(x)): 
    ny[i] = np.interp(nx, x[i], y[i]) 

Я использую метод слияния вместо метода двоичного поиска, потому что nx находится в порядке, а длина nx примерно такая же, как x.

• interp() использование бинарного поиска, время сложность O(len(nx)*log2(len(x))
• метод слияния: время сложность O(len(nx) + len(x))

Вот код Numba:

import numba 

@numba.jit("f8[::1](f8[::1], f8[::1], f8[::1], f8[::1])") 
def interp2(x, xp, fp, f): 
    n = len(x) 
    n2 = len(xp) 
    j = 0 
    i = 0 
    while x[i] <= xp[0]: 
     f[i] = fp[0] 
     i += 1 

    slope = (fp[j+1] - fp[j])/(xp[j+1] - xp[j])   
    while i < n: 
     if x[i] >= xp[j] and x[i] < xp[j+1]: 
      f[i] = slope*(x[i] - xp[j]) + fp[j] 
      i += 1 
      continue 
     j += 1 
     if j + 1 == n2: 
      break 
     slope = (fp[j+1] - fp[j])/(xp[j+1] - xp[j]) 

    while i < n: 
     f[i] = fp[n2-1] 
     i += 1 

@numba.jit("f8[:, ::1](f8[::1], f8[:, ::1], f8[:, ::1])") 
def multi_interp(x, xp, fp): 
    nrows = xp.shape[0] 
    f = np.empty((nrows, x.shape[0])) 
    for i in range(nrows): 
     interp2(x, xp[i, :], fp[i, :], f[i, :]) 
    return f 

Затем вызовите Numba функция:

%%time 
ny2 = multi_interp(nx, x, y) 

Чтобы проверить результат:

np.allclose(ny, ny2) 

На моем компьютере, время:

python version: 3.41 s 
numba version: 1.04 s 

Этот метод нужен массив, что последняя ось является осью быть interp().