存储高效存储大距离矩阵

时间:2018-05-04 00:14:17

标签: python numpy

我必须创建一个数据结构来存储从一个非常大的2d坐标数组中每个点到每个其他点的距离。对于小型阵列来说很容易实现,但超过大约50,000个点我开始遇到内存问题 - 这并不奇怪,因为我正在创建一个n x n矩阵。

这是一个很好的简单示例:

import numpy as np
from scipy.spatial import distance 

n = 2000
arr = np.random.rand(n,2)
d = distance.cdist(arr,arr)

cdist速度很快,但存储效率低,因为矩阵是对角镜像的(例如d[i][j] == d[j][i])。我可以使用np.triu(d)转换为上三角形,但生成的方形矩阵仍然使用相同的内存。我也不需要超出某个截止点的距离,因此这可能会有所帮助。下一步是转换为稀疏矩阵以节省内存:

from scipy import sparse

max_dist = 5
dist = np.array([[0,1,3,6], [1,0,8,7], [3,8,0,4], [6,7,4,0]])
print dist

array([[0, 1, 3, 6],
       [1, 0, 8, 7],
       [3, 8, 0, 4],
       [6, 7, 4, 0]])

dist[dist>=max_dist] = 0
dist = np.triu(dist)
print dist

array([[0, 1, 3, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 4],
       [0, 0, 0, 0]])

sdist = sparse.lil_matrix(dist)
print sdist

(0, 1)        1
(2, 3)        4
(0, 2)        3

对于非常大的数据集,问题是快速进入稀疏矩阵。重申一下,使用cdist制作方阵是我所知道的计算点之间距离的最快方法,但中间方阵矩阵耗尽内存。我可以把它分解成更易于处理的行块,但随后会减慢很多。我觉得我错过了一些从cdist直接转到稀疏矩阵的简单方法。

1 个答案:

答案 0 :(得分:2)

以下是使用KDTree

的方法
>>> import numpy as np
>>> from scipy import sparse
>>> from scipy.spatial import cKDTree as KDTree
>>> 
# mock data
>>> a = np.random.random((50000, 2))
>>> 
# make tree
>>> A = KDTree(a)
>>> 
# list all pairs within 0.05 of each other in 2-norm
# format: (i, j, v) - i, j are indices, v is distance
>>> D = A.sparse_distance_matrix(A, 0.05, p=2.0, output_type='ndarray')
>>> 
# only keep upper triangle
>>> DU = D[D['i'] < D['j']]
>>> 
# make sparse matrix
>>> result = sparse.coo_matrix((DU['v'], (DU['i'], DU['j'])), (50000, 50000))
>>> result
<50000x50000 sparse matrix of type '<class 'numpy.float64'>'
        with 9412560 stored elements in COOrdinate format>