加速 python 3.5 循环，使其以尽可能快的速度运行它

Question

我需要计算海量数据中 2 个 xyz 点之间的距离(100 Gb，大约 20 个 trylion 点(。我正在尝试加快这个循环。我创建了KDtree，添加了并行计算，将数组拆分为较小的部分。所以我想剩下的加速就是这个循环。我的纯 python 计算时间大约花了 10 小时 42 分钟。添加 numpy 将时间减少到 5 小时 34 分钟。添加 numba 将其速度提高到 4h 15 分钟。但它仍然不够快。我听说Cython是python计算的最快方法，但我没有任何c经验，我不知道如何将我的函数转换为cython代码。 如何使用cython或任何其他方式使此循环运行得更快？

def controller(point_array, las_point_array):  
    empty = []

    tree = spatial.cKDTree(point_array, leafsize=1000, copy_data = True)   
    empty = __pure_calc(las_point_array, point_array, empty, tree)  
    return ptList   
#############################################################################################
@autojit
def __pure_calc(las_point_array, point_array, empty, tree):
    for i in las_point_array:
            p = tree.query(i)   
            euc_dist = math.sqrt(np.sum((point_array[p[1]]-i)**2))  
            ##add one row at a time to empty list
            empty.append([i[0], i[1], i[2], euc_dist, point_array[p[1]][0], point_array[p[1]][1], point_array[p[1]][2]]) 
    return empty

我附上示例数据进行测试：

样本

Answer 1

你的函数构建了一个列表(closestPt (，最终看起来像这样：

[
    [i0[0], i0[1], i0[2], distM0],
    [i1[0], i1[1], i1[2], distM1],
    ...
]

您应该做的第一件事是将整个结果预分配为 NumPy 数组 ( np.empty() (，并一次写入一行。 这将避免大量的内存分配。 然后你会注意到，你可以将sqrt()推迟到最后，并在循环完成后在distM列上运行它。

如果您发布包含随机/样本输入数据的完整工作测试工具，则可能会有更多的优化机会。

Answer 2

关键是要尽可能多地利用矢量化函数，因为循环中对纯python函数的任何调用或多或少都会使autojit毫无意义(瓶颈将是纯函数调用(。我注意到查询函数是可矢量化的，欧几里得距离计算也是可矢量化的。我不确定控制器函数中的 ptList 变量是什么(该示例有点错误(，但假设它是 jit 函数的输出，或者接近它，您应该能够执行以下操作：

def controller(point_array, las_point_array):
    tree = spatial.cKDTree(point_array, leafsize=1000, copy_data = True)
    distances, pt_idx = tree.query(las_point_array)
    nearest_pts = point_array[pt_idx]
    euc_distances = np.sqrt((nearest_pts - las_point_array).sum(axis=1) ** 2)
    result = np.vstack((las_point_array.T, euc_distances.T, nearest_pts.T)).T
    return result