优化Python距离计算,同时考虑周期性边界条件

Question

我编写了一个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案例做相同的事情？有没有更好的选择？

Answer 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.

Answer 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 完成的。