計算 NumPy 坐標數組之間距離的三角矩陣

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 的pdist方法來獲得這些距離。 所以，我們只需要初始化輸出數組，然后用這些距離設置上三角值

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