¿Cuál es la mejor manera de calcular la traza de un producto matricial en números?

Si tengo una gran cantidad de matrices A y B , entonces puedo calcular la traza de su producto matricial con:

 tr = numpy.linalg.trace(A.dot(B)) 

Sin embargo, la multiplicación de la matriz A.dot(B) calcula innecesariamente todas las entradas fuera de la diagonal en el producto de la matriz, cuando solo se utilizan los elementos de la diagonal en la traza. En su lugar, podría hacer algo como:

 tr = 0.0 for i in range(n): tr += A[i, :].dot(B[:, i]) 

pero esto realiza el bucle en el código de Python y no es tan obvio como numpy.linalg.trace .

¿Hay una mejor manera de calcular la traza de un producto matricial de matrices numpy? ¿Cuál es la forma más rápida o más idiomática de hacer esto?

Puede mejorar la solución de @ Bill reduciendo el almacenamiento intermedio a los elementos diagonales solamente:

 from numpy.core.umath_tests import inner1d m, n = 1000, 500 a = np.random.rand(m, n) b = np.random.rand(n, m) # They all should give the same result print np.trace(a.dot(b)) print np.sum(a*bT) print np.sum(inner1d(a, bT)) %timeit np.trace(a.dot(b)) 10 loops, best of 3: 34.7 ms per loop %timeit np.sum(a*bT) 100 loops, best of 3: 4.85 ms per loop %timeit np.sum(inner1d(a, bT)) 1000 loops, best of 3: 1.83 ms per loop 

Otra opción es usar np.einsum y no tener ningún almacenamiento intermedio explícito:

 # Will print the same as the others: print np.einsum('ij,ji->', a, b) 

En mi sistema funciona un poco más lento que usando inner1d , pero puede que no sea inner1d para todos los sistemas, vea esta pregunta :

 %timeit np.einsum('ij,ji->', a, b) 100 loops, best of 3: 1.91 ms per loop 

Desde wikipedia puede calcular la traza utilizando el producto hammard (multiplicación de elementos):

 # Tr(AB) tr = (A*BT).sum() 

Creo que esto es menos cálculo que hacer numpy.trace(A.dot(B)) .

Editar:

Corrió algunos temporizadores. De esta manera es mucho más rápido que usar numpy.trace .

 In [37]: timeit("np.trace(A.dot(B))", setup="""import numpy as np; A, B = np.random.rand(1000,1000), np.random.rand(1000,1000)""", number=100) Out[38]: 8.6434469223022461 In [39]: timeit("(A*BT).sum()", setup="""import numpy as np; A, B = np.random.rand(1000,1000), np.random.rand(1000,1000)""", number=100) Out[40]: 0.5516049861907959 

Tenga en cuenta que una ligera variante es tomar el producto puntual de las matrices vecinas. En Python, la vectorización se realiza utilizando .flatten('F') . Es un poco más lento que tomar la sum del producto Hadamard en mi computadora, por lo que es una solución peor que la de Wflynny, pero creo que es algo interesante, ya que, en mi opinión, puede ser más intuitivo en algunas situaciones. Por ejemplo, personalmente encuentro que para la distribución normal de la matriz, la solución vectorizada es más fácil de entender.

Comparación de velocidad, en mi sistema:

 import numpy as np import time N = 1000 np.random.seed(123) A = np.random.randn(N, N) B = np.random.randn(N, N) tart = time.time() for i in range(10): C = np.trace(A.dot(B)) print(time.time() - start, C) start = time.time() for i in range(10): C = A.flatten('F').dot(BTflatten('F')) print(time.time() - start, C) start = time.time() for i in range(10): C = (AT * B).sum() print(time.time() - start, C) start = time.time() for i in range(10): C = (A * BT).sum() print(time.time() - start, C) 

Resultado:

 6.246593236923218 -629.370798672 0.06539678573608398 -629.370798672 0.057890892028808594 -629.370798672 0.05709719657897949 -629.370798672