Código de vectorización para calcular (al cuadrado) Mahalanobis Distiance

EDIT 2: esta publicación parece haberse trasladado de CrossValidated a StackOverflow debido a que se trata principalmente de progtwigción, pero eso significa que MathJax ya no funciona. Esperemos que esto todavía sea legible.

Digamos que quiero calcular la distancia de Mahalanobis al cuadrado entre dos vectores x y y con la matriz de covarianza S Esta es una función bastante simple definida por

 M2(x, y; S) = (x - y)^T * S^-1 * (x - y) 

Con el paquete numpy de python puedo hacer esto como

 # x, y = numpy.ndarray of shape (n,) # s_inv = numpy.ndarray of shape (n, n) diff = x - y d2 = diff.T.dot(s_inv).dot(diff) 

o en R como

 diff <- x - y d2 <- t(diff) %*% s_inv %*% diff 

En mi caso, sin embargo, me dan

  • m por n matriz X
  • vector tridimensional mu
  • n por n covarianza matriz S

y querer encontrar el vector tridimensional d tal que

 d_i = M2(x_i, mu; S) ( i = 1 .. m ) 

donde x_i es la i ª fila de X

Esto no es difícil de lograr usando un simple bucle en python:

 d = numpy.zeros((m,)) for i in range(m): diff = x[i,:] - mu d[i] = diff.T.dot(s_inv).dot(diff) 

Por supuesto, dado que el bucle externo está ocurriendo en python en lugar de en el código nativo en la biblioteca numpy significa que no es tan rápido como podría ser. $ n $ y $ m $ son aproximadamente 3-4 y varios cientos de miles, respectivamente, y estoy haciendo esto a menudo en un progtwig interactivo, por lo que una aceleración sería muy útil.

Matemáticamente, la única forma en que he podido formular esto usando operaciones matriciales básicas es

 d = diag( X' * S^-1 * X'^T ) 

dónde

  x'_i = x_i - mu 

la cual es simple para escribir una versión vectorizada de, pero desafortunadamente esto es superado por la ineficiencia de calcular una matriz de elementos de más de 10 mil millones y solo tomar la diagonal … Creo que esta operación debería ser fácilmente expresable usando la notación de Einstein, y por lo tanto podría ser evaluado rápidamente con la función einsum , pero ni siquiera he empezado a descubrir cómo funciona esa magia negra.

Por lo tanto, me gustaría saber: ¿existe alguna forma mejor de formular matemáticamente esta operación (en términos de operaciones matriciales simples) o alguien podría sugerir algún código vectorizado (python o R) que lo haga de manera eficiente?

PREGUNTA DE BONIFICACIÓN, para los valientes.

Realmente no quiero hacer esto una vez, quiero hacerlo k ~ 100 veces. Dado:

  • m por n matriz X

  • k por n matriz U

  • El conjunto de matrices de covarianza de n por n denotó S_j ( j = 1..k )

Encuentre la matriz D m por k tal que

 D_i,j = M(x_i, u_j; S_j) 

Donde i = 1..m , j = 1..k , x_i es la fila i th de X y u_j es la fila j th de U

Es decir, vectorizar el siguiente código:

 # s_inv is (kxnxn) array containing "stacked" inverses # of covariance matrices d = numpy.zeros( (m, k) ) for j in range(k): for i in range(m): diff = x[i, :] - u[j, :] d[i, j] = diff.T.dot(s_inv[j, :, :]).dot(diff) 

En primer lugar, parece que tal vez estás obteniendo S y luego invirtiéndolo. No deberías hacer eso; Es lento y numéricamente inexacto. En su lugar, debes obtener el factor de Cholesky L de S para que S = LL ^ T; entonces

 M^2(x, y; LL^T) = (x - y)^T (LL^T)^-1 (x - y) = (x - y)^TL^-TL^-1 (x - y) = || L^-1 (x - y) ||^2, 

y dado que L es triangular, L ^ -1 (x – y) se puede calcular de manera eficiente.

Resulta que, scipy.linalg.solve_triangular hará un montón de estos a la vez si lo remodelas correctamente:

 L = np.linalg.cholesky(S) y = scipy.linalg.solve_triangular(L, (X - mu[np.newaxis]).T, lower=True) d = np.einsum('ij,ij->j', y, y) 

Dividiéndolo un poco, y[i, j] es el componente i de L ^ -1 (X_j – \ mu). La llamada de einsum entonces hace

 d_j = \sum_i y_{ij} y_{ij} = \sum_i y_{ij}^2 = || y_j ||^2, 

como necesitamos

Desafortunadamente, solve_triangular no se vectorizará a través de su primer argumento, por lo que probablemente debería hacer un bucle allí. Si k es solo alrededor de 100, eso no será un problema significativo.


Si realmente le dan S ^ -1 en lugar de S, entonces puede hacerlo con einsum más directamente. Dado que S es bastante pequeño en su caso, también es posible que invertir la matriz y luego hacer esto sea más rápido. Sin embargo, tan pronto como n es un tamaño no trivial, estás descartando mucha precisión numérica al hacer esto.

Para averiguar qué hacer con einsum, escriba todo en términos de componentes. Iré directamente al caso de bonificación, escribiendo S_j ^ -1 = T_j para mayor comodidad de notación:

 D_{ij} = M^2(x_i, u_j; S_j) = (x_i - u_j)^T T_j (x_i - u_j) = \sum_k (x_i - u_j)_k ( T_j (x_i - u_j) )_k = \sum_k (x_i - u_j)_k \sum_l (T_j)_{kl} (x_i - u_j)_l = \sum_{kl} (X_{ik} - U_{jk}) (T_j)_{kl} (X_{il} - U_{jl}) 

Entonces, si hacemos matrices X de forma (m, n) , U de forma (k, n) y T de forma (k, n, n) , entonces podemos escribir esto como

 diff = X[np.newaxis, :, :] - U[:, np.newaxis, :] D = np.einsum('jik,jkl,jil->ij', diff, T, diff) 

donde diff[j, i, k] = X_[i, k] - U[j, k] .

Dougal le dio a esta una respuesta excelente y detallada, pero pensé que compartiría una pequeña modificación que encontré que aumenta la eficiencia en caso de que alguien más esté tratando de implementar esto. Directo al grano:

El método de Dougal fue el siguiente:

 def mahalanobis2(X, mu, sigma): L = np.linalg.cholesky(sigma) y = scipy.linalg.solve_triangular(L, (X - mu[np.newaxis,:]).T, lower=True) return np.einsum('ij,ij->j', y, y) 

Una variante matemáticamente equivalente que probé es

 def mahalanobis2_2(X, mu, sigma): # Cholesky decomposition of inverse of covariance matrix # (Doing this in either order should be equivalent) linv = np.linalg.cholesky(np.linalg.inv(sigma)) # Just do regular matrix multiplication with this matrix y = (X - mu[np.newaxis,:]).dot(linv) # Same as above, but note different index at end because the matrix # y is transposed here compared to above return np.einsum('ij,ij->i', y, y) 

Corrió ambas versiones de forma directa 20x usando entradas aleatorias idénticas y registró los tiempos (en milisegundos). Para X como una matriz de 1,000,000 x 3 (mu y sigma 3 y 3×3) obtengo:

 Method 1 (min/max/avg): 30/62/49 Method 2 (min/max/avg): 30/47/37 

Eso es aproximadamente un 30% de aceleración para la segunda versión. Lo haré principalmente en 3 o 4 dimensiones, pero para ver cómo se amplió, probé X como 1,000,000 x 100 y obtuve:

 Method 1 (min/max/avg): 970/1134/1043 Method 2 (min/max/avg): 776/907/837 

que se trata de la misma mejora.


Mencioné esto en un comentario sobre la respuesta de Dougal pero agregando aquí para mayor visibilidad:

El primer par de métodos de arriba toma un solo punto central mu y una matriz de covarianza sigma y calcula la distancia al cuadrado de Mahalanobis a cada fila de X. Mi pregunta adicional era hacerlo varias veces con muchos conjuntos de mu y sigma y generar una salida bidimensional. matriz. El conjunto de métodos anterior se puede utilizar para lograr esto con un simple bucle for, pero Dougal también publicó un ejemplo más inteligente utilizando einsum .

Decidí comparar estos métodos entre sí usándolos para resolver el siguiente problema: Dada distribuciones normales dd- k (con centros almacenados en filas de k por d matriz U y matrices de covarianza en las dos últimas dimensiones de k por d por d matriz S ), encuentre la densidad en los n puntos almacenados en filas de n por d matriz X

La densidad de una distribución normal multivariable es una función de la distancia al cuadrado de Mahalanobis del punto a la media. Scipy tiene una implementación de esto como scipy.stats.multivariate_normal.pdf para usar como referencia. Corrí los tres métodos uno contra otro 10x usando parámetros aleatorios idénticos cada vez, con d=3, k=96, n=5e5 . Aquí están los resultados, en puntos / seg:

 [Method]: (min/max/avg) Scipy: 1.18e5/1.29e5/1.22e5 Fancy 1: 1.41e5/1.53e5/1.48e5 Fancy 2: 8.69e4/9.73e4/9.03e4 Fancy 2 (cheating version): 8.61e4/9.88e4/9.04e4 

donde Fancy 1 es el mejor de los dos métodos anteriores y Fancy2 es la segunda solución de Dougal. Como el Fancy 2 necesita calcular las inversas de todas las matrices de covarianza, también probé una “versión engañosa” donde se pasaron estas como un parámetro, pero parece que eso no hizo una diferencia. Había planeado incluir la implementación no vectorizada, pero eso fue tan lento que hubiera tomado todo el día.

Lo que podemos quitar de esto es que usar el primer método de Dougal es aproximadamente un 20% más rápido de lo que Scipy lo hace. Desafortunadamente, a pesar de su inteligencia, el segundo método es tan solo un 60% más rápido que el primero. Probablemente hay otras optimizaciones que se pueden hacer, pero esto ya es lo suficientemente rápido para mí.

También probé cómo esta escalado con mayor dimensionalidad. Con d=100, k=96, n=1e4 :

 Scipy: 7.81e3/7.91e3/7.86e3 Fancy 1: 1.03e4/1.15e4/1.08e4 Fancy 2: 3.75e3/4.10e3/3.95e3 Fancy 2 (cheating version): 3.58e3/4.09e3/3.85e3 

Fancy 1 parece tener una ventaja aún mayor esta vez. También vale la pena señalar que Scipy lanzó un LinAlgError 8/10 veces, probablemente debido a que algunas de mis matrices de covarianza de 100×100 generadas aleatoriamente eran casi singulares (lo que puede significar que los otros dos métodos no son tan numéricamente estables, en realidad no verifiqué los resultados). ).