Rotación rápida del tensor con NumPy

En el corazón de una aplicación (escrita en Python y usando NumPy ) necesito rotar un tensor de cuarto orden. En realidad, necesito girar muchos tensores muchas veces y este es mi cuello de botella. Mi implementación ingenua (abajo) que involucra ocho bucles nesteds parece ser bastante lenta, pero no veo una manera de aprovechar las operaciones matriciales de NumPy y, con suerte, acelerar las cosas. Tengo la sensación de que debería estar usando np.tensordot , pero no veo cómo.

Matemáticamente, los elementos del tensor girado, T ‘, están dados por: T’ ijkl = Σ g ia g jb g kc g ld T abcd con la sum sobre los índices repetidos en el lado derecho. T y Tprime son matrices de 3 * 3 * 3 * 3 NumPy y la matriz de rotación g es una matriz de 3 * 3 NumPy. Mi implementación lenta (tomando ~ 0.04 segundos por llamada) está abajo.

 #!/usr/bin/env python import numpy as np def rotT(T, g): Tprime = np.zeros((3,3,3,3)) for i in range(3): for j in range(3): for k in range(3): for l in range(3): for ii in range(3): for jj in range(3): for kk in range(3): for ll in range(3): gg = g[ii,i]*g[jj,j]*g[kk,k]*g[ll,l] Tprime[i,j,k,l] = Tprime[i,j,k,l] + \ gg*T[ii,jj,kk,ll] return Tprime if __name__ == "__main__": T = np.array([[[[ 4.66533067e+01, 5.84985000e-02, -5.37671310e-01], [ 5.84985000e-02, 1.56722231e+01, 2.32831900e-02], [ -5.37671310e-01, 2.32831900e-02, 1.33399259e+01]], [[ 4.60051700e-02, 1.54658176e+01, 2.19568200e-02], [ 1.54658176e+01, -5.18223500e-02, -1.52814920e-01], [ 2.19568200e-02, -1.52814920e-01, -2.43874100e-02]], [[ -5.35577630e-01, 1.95558600e-02, 1.31108757e+01], [ 1.95558600e-02, -1.51342210e-01, -6.67615000e-03], [ 1.31108757e+01, -6.67615000e-03, 6.90486240e-01]]], [[[ 4.60051700e-02, 1.54658176e+01, 2.19568200e-02], [ 1.54658176e+01, -5.18223500e-02, -1.52814920e-01], [ 2.19568200e-02, -1.52814920e-01, -2.43874100e-02]], [[ 1.57414726e+01, -3.86167500e-02, -1.55971950e-01], [ -3.86167500e-02, 4.65601977e+01, -3.57741000e-02], [ -1.55971950e-01, -3.57741000e-02, 1.34215636e+01]], [[ 2.58256300e-02, -1.49072770e-01, -7.38843000e-03], [ -1.49072770e-01, -3.63410500e-02, 1.32039847e+01], [ -7.38843000e-03, 1.32039847e+01, 1.38172700e-02]]], [[[ -5.35577630e-01, 1.95558600e-02, 1.31108757e+01], [ 1.95558600e-02, -1.51342210e-01, -6.67615000e-03], [ 1.31108757e+01, -6.67615000e-03, 6.90486240e-01]], [[ 2.58256300e-02, -1.49072770e-01, -7.38843000e-03], [ -1.49072770e-01, -3.63410500e-02, 1.32039847e+01], [ -7.38843000e-03, 1.32039847e+01, 1.38172700e-02]], [[ 1.33639532e+01, -1.26331100e-02, 6.84650400e-01], [ -1.26331100e-02, 1.34222177e+01, 1.67851800e-02], [ 6.84650400e-01, 1.67851800e-02, 4.89151396e+01]]]]) g = np.array([[ 0.79389393, 0.54184237, 0.27593346], [-0.59925749, 0.62028664, 0.50609776], [ 0.10306737, -0.56714313, 0.8171449 ]]) for i in range(100): Tprime = rotT(T,g) 

¿Hay alguna manera de hacer que esto vaya más rápido? Hacer que el código se generalice a otros rangos de tensor sería útil, pero es menos importante.

Para usar tensordot , calcule el producto externo de los tensores g :

 def rotT(T, g): gg = np.outer(g, g) gggg = np.outer(gg, gg).reshape(4 * g.shape) axes = ((0, 2, 4, 6), (0, 1, 2, 3)) return np.tensordot(gggg, T, axes) 

En mi sistema, esto es siete veces más rápido que la solución de Sven. Si el tensor g no cambia con frecuencia, también puede almacenar el tensor gggg . Si haces esto y enciendes algunas tensordot ( tensordot código tensordot , sin controles, sin formas genéricas), aún puedes hacerlo dos veces más rápido:

 def rotT(T, gggg): return np.dot(gggg.transpose((1, 3, 5, 7, 0, 2, 4, 6)).reshape((81, 81)), T.reshape(81, 1)).reshape((3, 3, 3, 3)) 

Resultados de timeit en mi portátil doméstico (500 iteraciones):

 Your original code: 19.471129179 Sven's code: 0.718412876129 My first code: 0.118047952652 My second code: 0.0690279006958 

Los números en mi máquina de trabajo son:

 Your original code: 9.77922987938 Sven's code: 0.137110948563 My first code: 0.0569641590118 My second code: 0.0308079719543 

Aquí está cómo hacerlo con un solo bucle de Python:

 def rotT(T, g): Tprime = T for i in range(4): slices = [None] * 4 slices[i] = slice(None) slices *= 2 Tprime = g[slices].T * Tprime return Tprime.sum(-1).sum(-1).sum(-1).sum(-1) 

Es cierto que esto es un poco difícil de entender a primera vista, pero es un poco más rápido 🙂

Gracias al arduo trabajo de M. Wiebe, la próxima versión de Numpy (que probablemente será de 1.6) lo hará aún más fácil:

 >>> Trot = np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T) 

Sin embargo, el enfoque de Philipp en este momento es 3 veces más rápido, pero quizás haya algo de margen de mejora. La diferencia de velocidad probablemente se debe principalmente a que tensord no puede desenrollar toda la operación como un producto de matriz única que se puede pasar a BLAS, y así evitar gran parte de la sobrecarga asociada con los arreglos pequeños, esto no es posible para el general de Einstein la sum, ya que no todas las operaciones que se pueden express en esta forma se resuelven en un solo producto de matriz.

Por curiosidad he comparado la implementación de Cython de un código ingenuo de la pregunta con el código numpy de la respuesta de @ Philipp . El código de Cython es cuatro veces más rápido en mi máquina:

 #cython: boundscheck=False, wraparound=False import numpy as np cimport numpy as np def rotT(np.ndarray[np.float64_t, ndim=4] T, np.ndarray[np.float64_t, ndim=2] g): cdef np.ndarray[np.float64_t, ndim=4] Tprime cdef Py_ssize_t i, j, k, l, ii, jj, kk, ll cdef np.float64_t gg Tprime = np.zeros((3,3,3,3), dtype=T.dtype) for i in range(3): for j in range(3): for k in range(3): for l in range(3): for ii in range(3): for jj in range(3): for kk in range(3): for ll in range(3): gg = g[ii,i]*g[jj,j]*g[kk,k]*g[ll,l] Tprime[i,j,k,l] = Tprime[i,j,k,l] + \ gg*T[ii,jj,kk,ll] return Tprime 

Pensé que podría aportar un punto de datos relativamente nuevo a estos puntos de referencia, utilizando parakeet , uno de los comstackdores JIT de numpy capacidad que surgieron en los últimos meses. (El otro que conozco es numba , pero no lo probé aquí).

Después de completar el proceso de instalación un tanto laberíntico para LLVM, puede decorar muchas funciones de numpy puros para (a menudo) acelerar su rendimiento:

 import numpy as np import parakeet @parakeet.jit def rotT(T, g): # ... 

Solo intenté aplicar el JIT al código de Andrew en la pregunta original, pero funciona bastante bien (> 10x de aceleración) por no tener que escribir ningún código nuevo en absoluto:

 andrew 10 loops, best of 3: 206 msec per loop andrew_jit 10 loops, best of 3: 13.3 msec per loop sven 100 loops, best of 3: 2.39 msec per loop philipp 1000 loops, best of 3: 0.879 msec per loop 

Para estos tiempos (en mi computadora portátil) ejecuté cada función diez veces, para darle al JIT la oportunidad de identificar y optimizar las rutas de código activo.

Curiosamente, las sugerencias de Sven y Philipp siguen siendo órdenes de magnitud más rápido!

Enfoque prospectivo y código de solución

Para la eficiencia de la memoria y, posteriormente, la eficiencia del rendimiento, podríamos utilizar la multiplicación de la matriz del tensor en pasos.

Para ilustrar los pasos involucrados, utilicemos la solución más simple con np.einsum by @pv. –

 np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T) 

Como se ve, estamos perdiendo la primera dimensión de g con tensor-multiplicación entre sus cuatro variantes y T

Hagamos esas reducciones de sum para las multiplicaciones de la matriz tensorial en pasos. Empecemos con la primera variante de g y T :

 p1 = np.einsum('abcd, ai->bcdi', T, g) 

Así, terminamos con un tensor de dimensiones como notación de cadena: bcdi . Los siguientes pasos implicarían la reducción de la sum de este tensor en comparación con el rest de las tres variantes g que se utilizan en la einsum original de einsum. Por lo tanto, la próxima reducción sería –

 p2 = np.einsum('bcdi, bj->cdij', p1, g) 

Como hemos visto, hemos perdido las dos primeras dimensiones con las notaciones de cadena: a , b . Continuamos dos pasos más para deshacernos de c y d también y nos quedaría con ijkl como resultado final, como tal:

 p3 = np.einsum('cdij, ck->dijk', p2, g) p4 = np.einsum('dijk, dl->ijkl', p3, g) 

Ahora, podríamos usar np.tensordot para estas reducciones de sum, que serían mucho más eficientes.

Implementación final

Por lo tanto, al pasar a np.tensordot , tendríamos la implementación final así:

 p1 = np.tensordot(T,g,axes=((0),(0))) p2 = np.tensordot(p1,g,axes=((0),(0))) p3 = np.tensordot(p2,g,axes=((0),(0))) out = np.tensordot(p3,g,axes=((0),(0))) 

Prueba de tiempo de ejecución

Probemos todos los enfoques basados ​​en NumPy publicados en otras publicaciones para resolver el problema de rendimiento.

Enfoques según funciones.

 def rotT_Philipp(T, g): # @Philipp's soln gg = np.outer(g, g) gggg = np.outer(gg, gg).reshape(4 * g.shape) axes = ((0, 2, 4, 6), (0, 1, 2, 3)) return np.tensordot(gggg, T, axes) def rotT_Sven(T, g): # @Sven Marnach's soln Tprime = T for i in range(4): slices = [None] * 4 slices[i] = slice(None) slices *= 2 Tprime = g[slices].T * Tprime return Tprime.sum(-1).sum(-1).sum(-1).sum(-1) def rotT_pv(T, g): # @pv.'s soln return np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T) def rotT_Divakar(T,g): # Posted in this post p1 = np.tensordot(T,g,axes=((0),(0))) p2 = np.tensordot(p1,g,axes=((0),(0))) p3 = np.tensordot(p2,g,axes=((0),(0))) p4 = np.tensordot(p3,g,axes=((0),(0))) return p4 

Tiempos con los tamaños originales del conjunto de datos –

 In [304]: # Setup inputs ...: T = np.random.rand(3,3,3,3) ...: g = np.random.rand(3,3) ...: In [305]: %timeit rotT(T, g) ...: %timeit rotT_pv(T, g) ...: %timeit rotT_Sven(T, g) ...: %timeit rotT_Philipp(T, g) ...: %timeit rotT_Divakar(T, g) ...: 100 loops, best of 3: 6.51 ms per loop 1000 loops, best of 3: 247 µs per loop 10000 loops, best of 3: 137 µs per loop 10000 loops, best of 3: 41.6 µs per loop 10000 loops, best of 3: 28.3 µs per loop In [306]: 6510.0/28.3 # Speedup with the proposed soln over original code Out[306]: 230.03533568904592 

Como se comentó al inicio de esta publicación, estamos tratando de lograr la eficiencia de la memoria y, por lo tanto, boost el rendimiento con ella. Probemos eso a medida que aumentamos los tamaños de los conjuntos de datos –

 In [307]: # Setup inputs ...: T = np.random.rand(5,5,5,5) ...: g = np.random.rand(5,5) ...: In [308]: %timeit rotT(T, g) ...: %timeit rotT_pv(T, g) ...: %timeit rotT_Sven(T, g) ...: %timeit rotT_Philipp(T, g) ...: %timeit rotT_Divakar(T, g) ...: 100 loops, best of 3: 6.54 ms per loop 100 loops, best of 3: 7.17 ms per loop 100 loops, best of 3: 2.7 ms per loop 1000 loops, best of 3: 1.47 ms per loop 10000 loops, best of 3: 39.9 µs per loop