如何对均为2D的x、xp、yp进行向量化线性插值？

问题场景 #

需要对二维数组x、xp、yp执行线性插值，每行的x元素需对应到同一行的xp和yp完成计算。现有循环实现可得到正确结果，但处理大数组时性能不足，需向量化优化方案：

示例输入代码：

import numpy as np

x = np.array([[2, 3], [1, 2]], dtype=float)
xp = np.array([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=float)
yp = np.array([[0, 1, 3, 6, 7], [0, 3, 5, 7, 9]], dtype=float)

现有循环实现：

result = np.array([np.interp(x[i], xp[i], yp[i]) for i in range(len(x))])

向量化实现方案 #

利用NumPy的向量化操作替代Python循环，通过np.searchsorted批量定位插值区间，再直接计算插值结果，全程无显式循环，性能大幅提升：

import numpy as np

x = np.array([[2, 3], [1, 2]], dtype=float)
xp = np.array([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=float)
yp = np.array([[0, 1, 3, 6, 7], [0, 3, 5, 7, 9]], dtype=float)

# 定位每个x元素在对应xp行中的区间左索引
indices = np.searchsorted(xp, x, side='right') - 1
# 处理边界：避免索引越界（适配x小于xp最小值或大于最大值的情况）
indices = np.clip(indices, 0, xp.shape[1] - 2)

# 提取区间两端的xp和yp对应值
x0 = np.take_along_axis(xp, indices, axis=1)
x1 = np.take_along_axis(xp, indices + 1, axis=1)
y0 = np.take_along_axis(yp, indices, axis=1)
y1 = np.take_along_axis(yp, indices + 1, axis=1)

# 计算线性插值结果
result = y0 + (x - x0) * (y1 - y0) / (x1 - x0)

print(result)
# 输出：[[3. 6.]
#        [3. 5.]]

方案说明 #

• 完全基于NumPy向量化操作，规避Python循环的性能开销，大数组场景下速度远优于循环实现
• np.searchsorted批量完成索引定位，np.take_along_axis确保行与行的对应关系不混乱
• 边界处理逻辑与np.interp行为一致：当x超出xp范围时，直接取yp的边界值

内容的提问来源于stack exchange，提问作者Stefan