在考虑周期性边界条件的同时优化Python距离计算 [英] Optimizing Python distance calculation while accounting for periodic boundary conditions
问题描述
我写了一个Python脚本来计算3D空间中两点之间的距离,同时考虑到周期性边界条件.问题是我需要对很多点进行此计算,并且计算速度很慢.这是我的功能.
I have written a Python script to calculate the distance between two points in 3D space while accounting for periodic boundary conditions. The problem is that I need to do this calculation for many, many points and the calculation is quite slow. Here is my function.
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案例
Recently I read that I can greatly increase performance by using list comprehension. The following works for the non-PBC case, but not for the PBC case
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案件,有什么办法可以做到这一点吗?有没有更好的选择?
Is there some way to do the equivalent of this for the PBC case? Is there an alternative that would work better?
推荐答案
您应该编写distance()
函数,以便可以对5711个点上的循环进行矢量化处理.以下实现接受点数组作为x0
或x1
参数:
You should write your distance()
function in a way that you can vectorise the loop over the 5711 points. The following implementation accepts an array of points as either the x0
or x1
parameter:
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
中的每个点之间的距离的数组.
The result is the array of distances between the points passed as second parameter to distance()
and each point in points
.
这篇关于在考虑周期性边界条件的同时优化Python距离计算的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!