单位单元格中的距离矩阵(考虑对称性)

Question

我面临一个计算大距离矩阵的问题。然而，这是一个特定的距离矩阵：它是一个单元中的点矩阵。这个函数得到分数坐标(在所有维度上都在0和1之间(，我想计算距离矩阵，考虑到在晶胞的每个邻居中都有一个相同的点副本，因此正确的距离可能是与副本的距离，而不是与晶胞内的另一点的距离。

你知道用scipy或numpy预编码的C库可以做些什么吗？我已经做了一个numba代码，它可以工作，但运行相当慢。这里我有一个13160点的列表，我想为其计算13160*13160距离矩阵，即包含173185600个元素。

原理是：对于每个坐标，计算第一个点与第二个点在单元内或在其两个邻域之一(前后(的平方分数距离。然后得到每个坐标的平方距离的最小值，并从笛卡尔坐标得到相应的欧氏距离。

目前所需时间为：40.82661843299866秒

你知道我是否可以通过任何方式让它运行得更快吗？还是我的数据集太大了，没有什么可做的了？

以下是代码：

def getDistInCell(fract, xyz, n_sg, a, b, c):                           #calculates the distance matrix accounting for symmetry   
dist = np.zeros((n_sg, n_sg))
for i in range(n_sg):
for j in range(n_sg):
#we evaluate the closest segment according to translation to neighbouring cells
diff_x = np.zeros((3))
diff_y = np.zeros((3))
diff_z = np.zeros((3))

diff_x[0] = (fract[i][0] - (fract[j][0] - 1))**2
diff_x[1] = (fract[i][0] - (fract[j][0]    ))**2
diff_x[2] = (fract[i][0] - (fract[j][0] + 1))**2

diff_y[0] = (fract[i][1] - (fract[j][1] - 1))**2
diff_y[1] = (fract[i][1] - (fract[j][1]    ))**2
diff_y[2] = (fract[i][1] - (fract[j][1] + 1))**2

diff_z[0] = (fract[i][2] - (fract[j][2] - 1))**2
diff_z[1] = (fract[i][2] - (fract[j][2]    ))**2
diff_z[2] = (fract[i][2] - (fract[j][2] + 1))**2

#get correct shifts
shx = np.argmin(diff_x) - 1
shy = np.argmin(diff_y) - 1
shz = np.argmin(diff_z) - 1

#compute cartesian distance
dist[i][j] = np.sqrt((xyz[i][0] - (xyz[j][0] + shx * a)) ** 2 + (xyz[i][1] - (xyz[j][1] + shy * b)) ** 2 + (xyz[i][2] - (xyz[j][2] + shz * c)) ** 2)

return dist

Answer 1

以下是基于BallTree 的解决方案的示意图

我创建随机点，113160

import numpy as np
n=13160 
np.random.seed(1)
points=np.random.uniform(size=(n,3))

创建镜像/对称性，例如

from itertools import product
def create_symmetries( points ):

symmetries = []

for sym in product([0,-1,1],[0,-1,1],[0,-1,1]):
new_symmetry = points.copy()

diff_x, diff_y, diff_z = sym
new_symmetry[:,0] = new_symmetry[:,0] + diff_x
new_symmetry[:,1] = new_symmetry[:,1] + diff_y
new_symmetry[:,2] = new_symmetry[:,2] + diff_z
symmetries.append(new_symmetry)

return symmetries

并创建包括对称性的更大的数据集；

all_symmetries = np.concatenate( create_symmetries(points) )

要获得最接近的对称，请使用k=2，因为最接近的是点本身，第二个最接近的则是最接近的任何对称(包括它自己的对称，所以要小心(

%%time 
import numpy as np
from sklearn.neighbors import BallTree

tree = BallTree(all_symmetries, leaf_size=15, metric='euclidean')
dist, idx = tree.query(points, k=2, return_distance=True)

这需要<500ms

CPU times: user 275 ms, sys: 2.77 ms, total: 277 ms
Wall time: 275 ms