Scipy 快速一维插值无循环

Question

我有两个二维数组，x(ni, nj)和y(ni, nj)，我需要在一个轴上进行插值。我的目标是对每个ni的最后一个轴进行插值。

我写了

import numpy as np
from scipy.interpolate import interp1d

z = np.asarray([200,300,400,500,600])
out = []
for i in range(ni):
    f = interp1d(x[i,:], y[i,:], kind='linear')
    out.append(f(z))
out = np.asarray(out)

不过，我觉得这个方法效率不高，如果数组的大小太大，速度会很慢。有没有更快的方法来对这样的多维数组进行插值？有没有办法在不使用循环的情况下进行线性和立方插值？谢谢。

Answer 1

一种优化方法是一次性分配结果数组，像这样：

import numpy as np
from scipy.interpolate import interp1d

z = np.asarray([200,300,400,500,600])
out = np.zeros( [ni, len(z)], dtype=np.float32 ) 
for i in range(ni):
    f = interp1d(x[i,:], y[i,:], kind='linear')
    out[i,:]=f(z)

这样做可以节省一些内存复制的操作，这些操作是在你实现的过程中发生的，特别是在调用 out.append(...) 时。

Answer 2

你提到的方法确实用了Python的循环，所以如果ni的值很大，速度会变得很慢。不过，如果你的ni值不大，就不用太担心这个问题。

我用下面的代码创建了一些示例输入数据：

def sample_data(n_i, n_j, z_shape) :
    x = np.random.rand(n_i, n_j) * 1000
    x.sort()
    x[:,0] = 0
    x[:, -1] = 1000
    y = np.random.rand(n_i, n_j)
    z = np.random.rand(*z_shape) * 1000
    return x, y, z

然后我用这两个线性插值的版本进行了测试：

def interp_1(x, y, z) :
    rows, cols = x.shape
    out = np.empty((rows,) + z.shape, dtype=y.dtype)
    for j in xrange(rows) :
        out[j] =interp1d(x[j], y[j], kind='linear', copy=False)(z)
    return out

def interp_2(x, y, z) :
    rows, cols = x.shape
    row_idx = np.arange(rows).reshape((rows,) + (1,) * z.ndim)
    col_idx = np.argmax(x.reshape(x.shape + (1,) * z.ndim) > z, axis=1) - 1
    ret = y[row_idx, col_idx + 1] - y[row_idx, col_idx]
    ret /= x[row_idx, col_idx + 1] - x[row_idx, col_idx]
    ret *= z - x[row_idx, col_idx]
    ret += y[row_idx, col_idx]
    return ret

interp_1是你代码的优化版本，参考了Dave的答案。interp_2是一个向量化的线性插值实现，完全避免了Python的循环。写这样的代码需要对numpy中的广播和索引有很好的理解，而且有些地方的优化程度可能不如interp1d。比如说，找到要插值的值所在的区间：interp1d会在找到区间后提前跳出循环，而上面的函数则是把值与所有区间进行比较。

所以，结果会很依赖于n_i和n_j的值，甚至还要看你要插值的值的数组z有多长。如果n_j小而n_i大，你应该期待interp_2会有优势；反之，如果n_i小而n_j大，则interp_1会更好。较小的z对interp_2有利，而较长的z则对interp_1更有优势。

我实际上对这两种方法进行了计时，使用了不同的n_i和n_j，对于形状为(5,)和(50,)的z，这里是结果图：

所以看起来对于形状为(5,)的z，当n_j < 1000时，你应该选择interp_2，而在其他情况下选择interp_1。不出所料，形状为(50,)的z的临界值不同，现在是n_j < 100。看起来如果n_j * len(z) > 5000，你应该坚持用你的代码，但如果不是，就可以改用像interp_2这样的代码，但这个结论有很多推测的成分！如果你想进一步自己实验，这里是我用来生成图表的代码。

n_s = np.logspace(1, 3.3, 25)
int_1 = np.empty((len(n_s),) * 2)
int_2 = np.empty((len(n_s),) * 2)
z_shape = (5,)

for i, n_i in enumerate(n_s) :
    print int(n_i)
    for j, n_j in enumerate(n_s) :
        x, y, z = sample_data(int(n_i), int(n_j), z_shape)
        int_1[i, j] = min(timeit.repeat('interp_1(x, y, z)',
                                        'from __main__ import interp_1, x, y, z',
                                        repeat=10, number=1))
        int_2[i, j] = min(timeit.repeat('interp_2(x, y, z)',
                                        'from __main__ import interp_2, x, y, z',
                                        repeat=10, number=1))

cs = plt.contour(n_s, n_s, np.transpose(int_1-int_2))
plt.clabel(cs, inline=1, fontsize=10)
plt.xlabel('n_i')
plt.ylabel('n_j')
plt.title('timeit(interp_2) - timeit(interp_1), z.shape=' + str(z_shape))
plt.show()