如何加快两个不同大小的数组之间的操作？答案

【问题标题】：How to speed up an operation between two arrays of different sizes?如何加快两个不同大小的数组之间的操作？
【发布时间】：2021-06-13 01:18:45
【问题描述】：

我有两个数组，它们是 3D 空间中的点列表。这些数组有不同的长度。

np.shape(arr1) == (34709, 3)
np.shape(arr2) == (4835053, 3)

在给定周期性边界条件的情况下，我有一个函数可以计算一个数组中的单个点与另一个数组中的所有点之间的毕达哥拉斯距离：

def pythag_periodic(array, point, dimensions):
    delta = np.abs(array - point)
    delta = np.where(delta > 0.5 * dimensions, delta - dimensions, delta)
    return np.sqrt((delta ** 2).sum(axis=-1))

我正在尝试将此操作应用于两个数组中的所有点。我有一个递归调用这个函数的循环，但速度非常慢。

for i in arr1:
    pp.append(pythag_periodic(arr2, i, dimensions))

任何关于如何加快速度的建议将不胜感激。

【问题讨论】：

我会看看多处理库 - 特别是 Pool。 docs.python.org/3/library/multiprocessing.html
您正在查看大小为34709 x 4835053 的最终矩阵。你确定你有记忆吗？
几件事情要检查，map reduce，PyCUDA，使用 C/C++ 绑定运行计算密集型部分。
您还可以更改计算欧几里德距离的行以使用 np.linalg.norm()
您可以将此数据分块并在单独的进程上运行并将结果合并为一个。

标签： python python-3.x numpy multidimensional-array numpy-ndarray

【解决方案1】：

您应该使用 numba ：https://numba.pydata.org/（披露：我不是作者）。它是一个在运行时将 Python 函数转换为优化的机器代码的库。因此，Numba 编译的 Python 数值算法可以接近 C 或 FORTRAN 的速度。

应用到您的代码非常简单。简而言之，导入库，然后使用装饰器。此外，您还有更多与您相关的选项，例如 Parallelize Your Algorithms（查看他们的网站）。例如：

from numba import jit


@jit(nopython=True)
def pythag_periodic(array, point, dimensions):
    delta = np.abs(array - point)
    delta = np.where(delta > 0.5 * dimensions, delta - dimensions, delta)
    return np.sqrt((delta ** 2).sum(axis=-1))

【讨论】：

请注意，仅通过 jitting 这个函数，Numba 不会提供很大的改进。这应该快得多，因为在 arr1 上的循环迭代也是 jitted 的。另请注意，Numba 甚至可以并行化这个包含循环（如果需要）。

【解决方案2】：

另一个很酷的选择是利用 Numpy 的 broadcasting（在索引数组时通过 None 关键字）和超级简洁的 einsum 函数来避免循环并分别同时执行求和和平方运算。

但是请注意，这种方法对于小矩阵来说会稍微慢一些，但是一旦你的元素大小超过 4000 个左右，它就会快得多。此外，请注意内存不足，因为矢量化有这个缺点（尽管您已经在代码中存储了 NxM 数组）。

import numpy as np

def pythag_periodic_vectorized(a1, a2):
    delta = np.abs(a1[:,None,:] - a2[None,...])
    delta = np.where(delta > 0.5 * a1.shape[1], delta - a1.shape[1], delta)

return np.sqrt(np.einsum("ijk,ijk->ij", delta, delta))

【讨论】：