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

Question

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

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))

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

Answer 1

你应该使用 numba ：https://numba.pydata.org/（披露：我不是作者）。它是一个在运行时将 Python 函数转换为优化机器代码的库。因此，Python 中 Numba 编译的数值算法可以接近 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))

Answer 2

另一个很酷的选择是利用 Numpy 的 broadcasting (via the None keyword when indexing the arrays) and the super neat einsum 函数来避免循环并分别同时执行求和和平方运算。

但是请注意，这种方法对于小矩阵来说稍微慢一些，但是一旦你的大小超过 4000 个元素，它就会快得多。另外，请注意运行超出 RAM，因为矢量化有这个缺点（尽管您已经在代码中存储了 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))

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

How to speed up an operation between two arrays of different sizes?

python

numpy

multidimensional-array

python-3.x

numpy-ndarray