在考虑周期性边界条件的同时优化 Python 距离计算答案

【问题标题】：Optimizing Python distance calculation while accounting for periodic boundary conditions在考虑周期性边界条件的同时优化 Python 距离计算
【发布时间】：2012-06-21 23:00:40
【问题描述】：

我编写了一个 Python 脚本来计算 3D 空间中两点之间的距离，同时考虑周期性边界条件。问题是我需要对很多很多点进行这个计算，而且计算速度很慢。这是我的功能。

def PBCdist(coord1,coord2,UC):
    dx = coord1[0] - coord2[0]
    if (abs(dx) > UC[0]*0.5):
       dx = UC[0] - dx
    dy = coord1[1] - coord2[1]
    if (abs(dy) > UC[1]*0.5):
       dy = UC[1] - dy
    dz = coord1[2] - coord2[2]
    if (abs(dz) > UC[2]*0.5):
       dz = UC[2] - dz
    dist = np.sqrt(dx**2 + dy**2 + dz**2)
    return dist

然后我就这样调用函数

for i, coord2 in enumerate(coordlist):
  if (PBCdist(coord1,coord2,UC) < radius):
      do something with i

最近我读到使用列表理解可以大大提高性能。以下适用于非 PBC 案例，但不适用于 PBC 案例

coord_indices = [i for i, y in enumerate([np.sqrt(np.sum((coord2-coord1)**2)) for coord2 in coordlist]) if y < radius]
for i in coord_indices:
   do something

对于 PBC 案，有什么方法可以做到这一点吗？有没有更好的替代方法？

【问题讨论】：

您正在使用 NumPy，因此您应该对循环进行矢量化以提高性能。 coordlist 的结构到底是什么？它应该是一个二维 NumPy 数组，以便能够使用 NumPy ufunc 优化循环。
coordlist 是一个形状约为 (5711,3) 的 numpy 数组。 coordlist 本身来自一个更大的列表，所以我有效地循环 coordlist 20,000 次，而 coordlist 的列表循环了大约 50 次......你明白了。
我在 NumPy 中查找了矢量化函数。文档说："The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop."
"Vectorising" 在 NumPy 的上下文中意味着您使用 NumPy ufunc 将循环移动，否则您必须在 Python 代码中执行 C 代码（就像我在回答中所做的那样）。如果您担心性能，我建议您阅读一下 NumPy。

标签： python optimization list-comprehension

【解决方案1】：

您应该编写您的 distance() 函数，以便您可以将循环矢量化到 5711 个点上。以下实现接受点数组作为x0 或x1 参数：

def distance(x0, x1, dimensions):
    delta = numpy.abs(x0 - x1)
    delta = numpy.where(delta > 0.5 * dimensions, delta - dimensions, delta)
    return numpy.sqrt((delta ** 2).sum(axis=-1))

例子：

>>> dimensions = numpy.array([3.0, 4.0, 5.0])
>>> points = numpy.array([[2.7, 1.5, 4.3], [1.2, 0.3, 4.2]])
>>> distance(points, [1.5, 2.0, 2.5], dimensions)
array([ 2.22036033,  2.42280829])

结果是作为第二个参数传递给distance() 的点与points 中的每个点之间的距离数组。

【讨论】：

这导致我的代码速度提高了大约 5 倍。谢谢！对于那些寻找替代方案的人来说，lazyr 的回答也表现得一样好。
采用规范后没有区别，但我喜欢通过将第 3 行替换为 delta = numpy.where(", delta - dimensions, ") 来获得正确的 delta 符号。另请注意，np.hypot 在避免溢出方面比 sqrt(sum(x**2))
@Patrick numpy.hypot() 仅适用于二维，而 OP 需要适用于三维点的代码。关于delta 的标志，好吧，如果您关心，请随时编辑帖子。 :)
我认为当点不在相邻的盒子中，但距离另一个盒子不止一个盒子时，这可能会失败。
@k1next 不确定你在说什么盒子。

【解决方案2】：

import numpy as np

bounds = np.array([10, 10, 10])
a = np.array([[0, 3, 9], [1, 1, 1]])
b = np.array([[2, 9, 1], [5, 6, 7]])

min_dists = np.min(np.dstack(((a - b) % bounds, (b - a) % bounds)), axis = 2)
dists = np.sqrt(np.sum(min_dists ** 2, axis = 1))

这里的a 和b 是您希望计算距离的向量列表，bounds 是空间的边界（所以这里所有三个维度都是从 0 到 10 然后换行）。它计算a[0]和b[0]、a[1]和b[1]之间的距离，等等。

我相信 numpy 专家可以做得更好，但这可能会比你正在做的快一个数量级，因为现在大部分工作都是用 C 完成的。

【讨论】：

我也尝试过这种方法。它还导致代码增强了大约 5 倍。不幸的是，我只能检查 1 个答案是否正确：/
@johnjax 对于它的价值，如果我站在你的立场上，我也会接受 Sven Marnach 的回答。它比我的更直接适用。

【解决方案3】：

我发现meshgrid 对于生成距离非常有用。例如：

import numpy as np
row_diff, col_diff = np.meshgrid(range(7), range(8))
radius_squared = (row_diff - x_coord)**2 + (col_diff - y_coord)**2

我现在有一个数组 (radius_squared)，其中每个条目都指定与数组位置 [x_coord, y_coord] 的距离的平方。

要循环数组，我可以执行以下操作：

row_diff, col_diff = np.meshgrid(range(7), range(8))
row_diff = np.abs(row_diff - x_coord)
row_circ_idx = np.where(row_diff > row_diff.shape[1] / 2)
row_diff[row_circ_idx] = (row_diff[row_circ_idx] - 
                         2 * (row_circ_idx + x_coord) + 
                         row_diff.shape[1])
row_diff = np.abs(row_diff)
col_diff = np.abs(col_diff - y_coord)
col_circ_idx = np.where(col_diff > col_diff.shape[0] / 2)
col_diff[row_circ_idx] = (row_diff[col_circ_idx] - 
                         2 * (col_circ_idx + y_coord) + 
                         col_diff.shape[0])
col_diff = np.abs(row_diff)
circular_radius_squared = (row_diff - x_coord)**2 + (col_diff - y_coord)**2

我现在已经用矢量数学循环了所有的数组距离。

【讨论】：