考虑周期边界条件的Python距离优化计算

2024-10-01 09:30:47 发布

您现在位置:Python中文网/ 问答频道 /正文

我编写了一个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

对于中国人民银行的案子,有什么方法可以做到这一点吗?有没有更好的选择?


Tags: inforifdistnpabssqrtuc