本文介绍如何将单点线性插值函数改造为支持 numpy 1d 数组输入的向量化版本，在保持 numba 加速优势的同时避免 python 循环开销，显著提升优化算法中高频插值的执行效率。
本文介绍如何将单点线性插值函数改造为支持 numpy 1d 数组输入的向量化版本，在保持 numba 加速优势的同时避免 python 循环开销，显著提升优化算法中高频插值的执行效率。
原始 calc 函数使用 @njit 装饰，已具备出色的单点插值性能，但其设计仅接受标量 x0，无法直接处理数组。若在 Python 层面用 for 循环逐个调用（如答案中所示），虽能工作，却丧失了向量化潜力——每次调用都触发一次 JIT 函数调度，且无法利用 CPU 向量指令与缓存局部性，性能远未达最优。
更高效的方案是将插值逻辑本身向量化，即改写 calc 函数以原生支持 x0 为 1D NumPy 数组（float64[:]）。Numba 完全支持此类数组操作，关键在于：
- 使用
np.empty_like(x0)预分配结果数组； - 利用
prange（并行循环）替代普通range，启用多线程加速； - 保持边界检查与分段线性逻辑不变，但对每个索引独立计算。
以下是优化后的完整实现：
import numpy as np
from numba import njit, prange
# 向量化插值函数（支持标量与1D数组）
@njit(parallel=True)
def calc_vectorized(x0, x, y):
n = len(x0)
result = np.empty(n, dtype=np.float64)
for i in prange(n): # 并行遍历输入数组
x_val = x0[i]
if x_val <= x[0]:
result[i] = y[0]
elif x_val >= x[-1]:
result[i] = y[-1]
else:
# 二分查找可进一步加速（适用于长表），此处保留线性搜索以保证简洁性
# 注意：x 必须严格递增，否则行为未定义
for j in range(len(x) - 1):
if x[j] <= x_val <= x[j + 1]:
x1, x2 = x[j], x[j + 1]
y1, y2 = y[j], y[j + 1]
result[i] = y1 + (y2 - y1) / (x2 - x1) * (x_val - x1)
break
return result
# 插值曲线封装（自动处理标量/数组输入）
WeirCurve = np.array([[749.81, 0], [749.9, 5], [750, 14.2],
[751, 226], [752, 556], [753, 923.2], [754, 1155.3]])
def WeirDischCurve(x):
x = np.asarray(x)
if x.ndim == 0: # 标量输入
return calc(x.item(), WeirCurve[:, 0], WeirCurve[:, 1])
elif x.ndim == 1: # 1D 数组输入
return calc_vectorized(x, WeirCurve[:, 0], WeirCurve[:, 1])
else:
raise ValueError("Only scalar or 1D array inputs are supported.")✅ 使用示例：
print(WeirDischCurve(751.65)) # 标量 → 440.5 print(WeirDischCurve([751.65, 752.5, 753.3])) # 向量 → [440.5, 739.6, 992.83]
⚠️ 关键注意事项：
-
输入
x（节点横坐标）必须严格单调递增，否则区间查找逻辑失效；建议在初始化时校验np.all(np.diff(x) > 0)。 - 若插值表非常长（>1000 点），可将内部循环替换为
numba.typed.List+ 二分搜索（numba.experimental.jitclass或手动实现），将单次查找复杂度从 O(n) 降至 O(log n)。 -
@njit(parallel=True)在多核 CPU 上可带来近线性加速比，但需确保x0长度足够大（通常 ≥ 1000）以摊销并行开销。 - 避免在
calc_vectorized中传入非编译类型（如 Python list），务必使用np.array(..., dtype=np.float64)。
该方案在不依赖 scipy.interpolate 的前提下，兼顾精度、速度与内存效率，特别适合嵌入数值优化内层循环的高性能场景。
