如何用NumPy高效将3D函数值映射填充到三维网格

核心优化思路 #

你现有代码的性能瓶颈来自四重循环的暴力区间匹配，时间复杂度高达O(n*N^3)。由于你使用的是固定步长的等距网格，可以直接通过算术运算向量化计算所有采样点对应的网格索引，时间复杂度直接降到O(n)，完全不需要循环。

优化后代码 #

import numpy as np

step = 0.1

# 计算各维度边界
xmin, ymin, zmin = data.min(axis=0)
xmax, ymax, zmax = data.max(axis=0)

# 计算各维度网格长度
nx = int((xmax - xmin)/step) + 1
ny = int((ymax - ymin)/step) + 1
nz = int((zmax - zmin)/step) + 1

# 初始化网格
M = np.zeros((nx, ny, nz))

# 向量化计算所有点对应的网格索引
i = np.floor((data[:, 0] - xmin) / step).astype(int)
j = np.floor((data[:, 1] - ymin) / step).astype(int)
k = np.floor((data[:, 2] - zmin) / step).astype(int)

# 防止浮点误差导致索引越界
i = np.clip(i, 0, nx - 1)
j = np.clip(j, 0, ny - 1)
k = np.clip(k, 0, nz - 1)

# 赋值到网格，若同个网格有多个点默认保留最后一个点的值
M[i, j, k] = F

# 如果同个网格多个点需要累加值，替换上面的赋值为下面的代码即可
# np.add.at(M, (i, j, k), F)

不等距网格适配 #

如果你后续需要使用非等距的自定义网格，可以把索引计算部分替换为searchsorted实现，同样不需要循环：

# 假设x/y/z是自定义的非等距网格点数组
i = np.searchsorted(x, data[:, 0]) - 1
j = np.searchsorted(y, data[:, 1]) - 1
k = np.searchsorted(z, data[:, 2]) - 1

# 做边界截断避免越界
i = np.clip(i, 0, len(x)-2)
j = np.clip(j, 0, len(y)-2)
k = np.clip(k, 0, len(z)-2)

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