Cómo hacer la distancia nD y los cálculos del vecino más cercano en matrices numpy

Esta pregunta pretende ser un objective duplicado canónico

Dados dos matrices X e Y de formas (i, n) y (j, n) , que representan listas de coordenadas n dimensionales,

 def test_data(n, i, j, r = 100): X = np.random.rand(i, n) * r - r / 2 Y = np.random.rand(j, n) * r - r / 2 return X, Y X, Y = test_data(3, 1000, 1000) 

¿Cuáles son las formas más rápidas de encontrar:

  1. La distancia D con forma (i,j) entre cada punto en X y cada punto en Y
  2. Los índices k_i y la distancia k_d de los k vecinos más cercanos contra todos los puntos en X para cada punto en Y
  3. Los índices r_i , r_j y la distancia r_d de cada punto en X dentro de la distancia r de cada punto j en Y

Teniendo en cuenta los siguientes conjuntos de restricciones:

  • Sólo usando numpy
  • Usando cualquier paquete de python

Incluyendo el caso especial:

  • Y es X

En todos los casos, la distancia significa principalmente la distancia euclidiana , pero no dude en resaltar los métodos que permiten otros cálculos de distancia.

1. Todas las distancias

  • solo usando numpy

El método ingenuo es:

 D = np.sqrt(np.sum((X[:, None, :] - Y[None, :, :])**2, axis = -1)) 

Sin embargo, esto requiere una gran cantidad de memoria creando una matriz intermedia con forma (i, j, n) y es muy lento

Sin embargo, gracias a un truco de @Divakar (paquete eucl_dist , wiki ), podemos usar un poco de álgebra y np.einsum para descomponer como tal: (X - Y)**2 = X**2 - 2*X*Y + Y**2

 D = np.sqrt( # (X - Y) ** 2 np.einsum('ij, ij ->i', X, X)[:, None] + # = X ** 2 \ np.einsum('ij, ij ->i', Y, Y) - # + Y ** 2 \ 2 * X.dot(YT)) # - 2 * X * Y 
  • Y es X

Similar a lo anterior:

 XX = np.einsum('ij, ij ->i', X, X) D = np.sqrt(XX[:, None] + XX - 2 * X.dot(XT)) 

Tenga en cuenta que la imprecisión de punto flotante puede hacer que los términos diagonales se desvíen ligeramente de cero con este método. Si necesita asegurarse de que son cero, deberá establecerlo explícitamente:

 np.einsum('ii->i', D)[:] = 0 
  • Cualquier paquete

scipy.spatial.distance.cdist es la función incorporada más intuitiva para esto, y mucho más rápida que la numpy

 from scipy.spatial.distance import cdist D = cdist(X, Y) 

cdist también puede manejar muchas, muchas medidas de distancia así como medidas de distancia definidas por el usuario (aunque no están optimizadas). Consulte la documentación vinculada anteriormente para más detalles.

  • Y es X

Para distancias scipy.spatial.distance.pdist , scipy.spatial.distance.pdist funciona de manera similar a cdist , pero devuelve una matriz de distancia condensada 1-D, ahorrando espacio en la matriz de distancia simétrica al tener solo cada término una vez. Puedes convertir esto en una matriz cuadrada usando squareform

 from scipy.spatial.distance import pdist, squareform D_cond = pdist(X) D = squareform(D_cond) 

2. K vecinos más cercanos (KNN)

  • Sólo usando numpy

Podríamos usar np.argpartition para obtener los índices k-nearest y usarlos para obtener los valores de distancia correspondientes. Entonces, con D como la matriz que contiene los valores de distancia obtenidos anteriormente, tendríamos –

 if k == 1: k_i = D.argmin(0) else: k_i = D.argpartition(k, axis = 0)[:k] k_d = np.take_along_axis(D, k_i, axis = 0) 

Sin embargo, podemos acelerar esto un poco si no tomamos las raíces cuadradas hasta que hayamos reducido nuestro conjunto de datos. np.sqrt es la parte más lenta de calcular la norma euclidiana, por lo que no queremos hacerlo hasta el final.

 D_sq = np.einsum('ij, ij ->i', X, X)[:, None] +\ np.einsum('ij, ij ->i', Y, Y) - 2 * X.dot(YT) if k == 1: k_i = D_sq.argmin(0) else: k_i = D_sq.argpartition(k, axis = 0)[:k] k_d = np.sqrt(np.take_along_axis(D_sq, k_i, axis = 0)) 

Ahora, np.argpartition realiza una partición indirecta y no necesariamente nos da los elementos ordenados y solo se asegura de que los primeros k elementos sean los más pequeños. Por lo tanto, para una salida ordenada, debemos usar argsort en la salida del paso anterior:

 sorted_idx = k_d.argsort(axis = 0) k_i_sorted = np.take_along_axis(k_i, sorted_idx, axis = 0) k_d_sorted = np.take_along_axis(dist, k_i_sorted, axis = 0) 
  • Cualquier paquete

KD-Tree es un método mucho más rápido para encontrar vecinos y distancias limitadas. El método más recomendado es usar scipy ‘s scipy.spatial.KDTree o scipy.spatial.cKDTree

 from scipy.spatial.distance import KDTree X_tree = KDTree(X) k_d, k_i = X_tree.query(Y, k = k) 

Desafortunadamente, la implementación de KDTree de scipy es lenta y tiene una tendencia a segfault para conjuntos de datos más grandes. Como lo señala @HansMusgrave aquí , pykdtree aumenta mucho el rendimiento, pero no es tan común como scipy y solo puede lidiar con la distancia euclidiana actualmente (mientras que KDTree en scipy puede manejar las p-normas de Minkowsi de cualquier orden)

  • Métricas arbitrarias

Un BallTree tiene propiedades algorítmicas similares a un KDTree. No tengo conocimiento de un BallTree paralelo / vectorizado / rápido en Python, pero al usar scipy aún podemos tener consultas KNN razonables para las métricas definidas por el usuario. Si están disponibles, las métricas integradas serán mucho más rápidas.

 def d(a, b): return max(np.abs(ab)) tree = sklearn.neighbors.BallTree(X, metric=d) k_d, k_i = tree.query(Y) 

Esta respuesta será incorrecta si d() no es una métrica . La única razón por la que un BallTree es más rápido que la fuerza bruta es porque las propiedades de una métrica le permiten descartar algunas soluciones. Para funciones verdaderamente arbitrarias, la fuerza bruta es realmente necesaria.

3. Radio de búsqueda

  • Sólo usando numpy

El método más simple es usar la indexación booleana:

 mask = D_sq < r**2 r_i, r_j = np.where(mask) r_d = np.sqrt(D_sq[mask]) 
  • Cualquier paquete

Similar a lo anterior, puedes usar scipy.spatial.KDTree.query_ball_point

 r_ij = X_tree.query_ball_point(Y, r = r) 

o scipy.spatial.KDTree.query_ball_tree

 Y_tree = KDTree(Y) r_ij = X_tree.query_ball_tree(Y_tree, r = r) 

Lamentablemente, r_ij termina siendo una lista de matrices de índices que son un poco difíciles de desentrañar para su uso posterior.

Mucho más fácil es usar el cKDTree de sparse_distance_matrix , que puede generar un coo_matrix

 from scipy.spatial.distance import cKDTree X_cTree = cKDTree(X) Y_cTree = cKDTree(Y) D_coo = X_cTree.sparse_distance_matrix(Y_cTree, r = r, output_type = `coo_matrix`) r_i = D_coo.row r_j = D_coo.column r_d = D_coo.data 

Este es un formato extraordinariamente flexible para la matriz de distancia, ya que se mantiene como una matriz real que (si se convierte a csr ) también se puede usar para muchas operaciones vectorizadas.