首页 文章

计算NumPy坐标数组之间距离的三角矩阵

提问于
浏览
2

我有一个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秒)?

1 回答

  • 2

    我们可以使用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)
    

    Runtime test

    功能:

    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
    

相关问题