Question

我有一个NumPy坐标数组。例如，我将使用此

In [1]: np.random.seed(123)
In [2]: coor = np.random.randint(10, size=12).reshape(-1,3)
In [3]: coor
Out[3]: array([[2, 2, 6],
               [1, 3, 9],
               [6, 1, 0],
               [1, 9, 0]])

我想要所有坐标之间距离的三角矩阵。一种简单的方法是在所有坐标上编码双循环

In [4]: n_coor = len(coor)
In [5]: dist = np.zeros((n_coor, n_coor))
In [6]: for j in xrange(n_coor):
            for k in xrange(j+1, n_coor):
                dist[j, k] = np.sqrt(np.sum((coor[j] - coor[k]) ** 2))

结果是距离的上三角矩阵

In [7]: dist
Out[7]: array([[  0.        ,   3.31662479,   7.28010989,   9.2736185 ],
               [  0.        ,   0.        ,  10.48808848,  10.81665383],
               [  0.        ,   0.        ,   0.        ,   9.43398113],
               [  0.        ,   0.        ,   0.        ,   0.        ]])

利用NumPy，我可以避免使用

进行循环

In [8]: dist = np.sqrt(((coor[:, None, :] - coor) ** 2).sum(-1))

但结果是整个矩阵

In [9]: dist
Out[9]: array([[  0.        ,   3.31662479,   7.28010989,   9.2736185 ],
               [  3.31662479,   0.        ,  10.48808848,  10.81665383],
               [  7.28010989,  10.48808848,   0.        ,   9.43398113],
               [  9.2736185 ,  10.81665383,   9.43398113,   0.        ]])

当我使用2048坐标（4秒而不是10秒）时，这一行版本大约需要一半的时间，但这需要两倍的计算才能获得对称矩阵。有没有办法调整一行命令只能获得三角矩阵（以及额外的2倍加速，即2秒）？

Answer 1

我们可以使用SciPy's pdist method来获取这些距离。因此，我们只需要初始化输出数组，然后使用这些距离设置上三角值

from scipy.spatial.distance import pdist

n_coor = len(coor)
dist = np.zeros((n_coor, n_coor))
row,col = np.triu_indices(n_coor,1)
dist[row,col] = pdist(coor)

或者，我们可以使用布尔索引来分配值，替换最后两行

dist[np.arange(n_coor)[:,None] < np.arange(n_coor)] = pdist(coor)

运行时测试

功能：

def subscripted_indexing(coor):
    n_coor = len(coor)
    dist = np.zeros((n_coor, n_coor))
    row,col = np.triu_indices(n_coor,1)
    dist[row,col] = pdist(coor)
    return dist

def boolean_indexing(coor):
    n_coor = len(coor)
    dist = np.zeros((n_coor, n_coor))
    r = np.arange(n_coor)
    dist[r[:,None] < r] = pdist(coor)
    return dist

时序：

In [110]: # Setup input array
     ...: coor = np.random.randint(0,10, (2048,3))

In [111]: %timeit subscripted_indexing(coor)
10 loops, best of 3: 91.4 ms per loop

In [112]: %timeit boolean_indexing(coor)
10 loops, best of 3: 47.8 ms per loop

计算NumPy坐标数组

1 个答案: