我编写了一个 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 案做同样的事情?有没有更好的替代方案?