Numpy带参数高效函数映射：高度图插值代码性能优化

问题背景 #

我尝试通过区域内若干点的高度值插值生成高度图，初始实现处理大尺寸图像（如800×600）耗时近一分钟，时间复杂度为O(n²)。目标是在不改变时间复杂度的前提下减少耗时，快速生成合适尺寸的图像。

初始实现代码 #

生成高度图的主逻辑：

map_ = np.zeros((img_width, img_height))

for x in range(img_width):
    for y in range(img_height):
        map_[x, y] = calculate_height(set(points.items()), x, y)

calculate_height函数实现：

def distance(x1, y1, x2, y2) -> float:
    return np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)


def calculate_height(points: set, x, y) -> float:
    total = 0
    dists = {}
    for pos, h in points:
        d = distance(pos[0], pos[1], x, y)
        if x == pos[0] and y == pos[1]:
            return h
        d = 1 / (d ** 2)
        dists[pos] = d
        total += d

    r = 0
    for pos, h in points:
        ratio = dists[pos] / total
        r += ratio * h

    return r

points的数据结构：

points: dict[tuple[int, int], float] = {
    (x: int, y: int): height: float,
    (x: int, y: int): height: float,
    (x: int, y: int): height: float
}
# 点的数量范围为 [3, +∞)

已完成的优化 #

优化1：避免重复构造集合 #

将set(points.items())移出循环，不再每次调用calculate_height时重复创建集合，性能大幅提升。

优化2：简化距离计算逻辑 #

修改calculate_height中的距离计算部分，直接使用距离平方的倒数，避免开根号操作，进一步降低耗时：

# calculate_height函数内的第一个循环修改为：
for pos, h in points:
    d2 = distance2(pos[0], pos[1], x, y)  # distance2返回(x1-x2)² + (y1-y2)²
    if x == pos[0] and y == pos[1]:
        return h
    d2 = d2 ** -1  # 1/(d²) 等价于 (d²)^-1
    dists[pos] = d2
    total += d2

优化后，200×200图像的处理时间从1.57秒降至0.76秒，800×600图像从近一分钟降至6.13秒。

更多优化建议 #

• 向量化计算替代循环：利用NumPy的广播机制，一次性计算所有目标点与散点的距离平方，避免Python层面的循环。提前将散点坐标和高度转为NumPy数组，对整个高度图的网格坐标进行批量计算，减少循环开销。
• 移除冗余字典操作：calculate_height中使用dists字典存储距离权重是冗余的，可在第一次循环时直接计算加权高度和，无需第二次遍历。示例：

  def calculate_height(pos_array, h_array, x, y):
      # 检查是否是散点本身
      mask = (pos_array[:,0] == x) & (pos_array[:,1] == y)
      if np.any(mask):
          return h_array[mask][0]
      # 计算距离平方与权重
      dx = pos_array[:,0] - x
      dy = pos_array[:,1] - y
      d2 = dx**2 + dy**2
      weights = 1 / d2
      total_weight = weights.sum()
      return (weights * h_array).sum() / total_weight
  

• JIT编译加速：用numba库的@njit装饰器修饰calculate_height函数，将Python代码编译为机器码执行，大幅提升循环密集型代码的速度。
• 并行分块处理：将大尺寸图像分成若干小块，利用multiprocessing或concurrent.futures实现多CPU核心并行计算，进一步缩短总耗时。

内容的提问来源于stack exchange，提问作者Ciro García