python - 3D 原点 x 的 3D 数组快速插值

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() 

在 numpy 中进行 200000 次插入:

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

我用的是merge方法而不是二分查找方法，因为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() 的轴。