Эффективное для памяти хранение больших дистанционных матриц

Question

Мне нужно создать структуру данных для хранения расстояний от каждой точки до каждой другой точки в очень большом массиве 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.

Answer 1

Вот как это сделать с 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>