Rotation rapide de tenseurs avec NumPy

Question

Au cœur d'une application (écrite en Python et utilisant NumPy ) J'ai besoin de faire tourner un tenseur d'ordre 4. En fait, j'ai besoin de faire tourner un grand nombre de tenseurs plusieurs fois et c'est là mon goulot d'étranglement. Mon implémentation naïve (ci-dessous) impliquant huit boucles imbriquées semble être assez lente, mais je ne vois pas de moyen de tirer parti des opérations matricielles de NumPy et, espérons-le, d'accélérer les choses. J'ai l'impression que je devrais utiliser np.tensordot mais je ne vois pas comment.

Mathématiquement, les éléments du tenseur tourné, T', sont donnés par : T' _ijkl \= Σ g _ia g _jb g _kc g _ld T _abcd avec la somme sur les indices répétés sur le côté droit. T et Tprime sont des tableaux NumPy 3*3*3*3 et la matrice de rotation g est un tableau NumPy 3*3. Mon implémentation lente (prenant ~0.04 secondes par appel) est ci-dessous.

#!/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)

Y a-t-il un moyen d'accélérer le processus ? Généraliser le code à d'autres rangs de tenseurs serait utile, mais c'est moins important.

Answer 1

Pour utiliser tensordot calculer le produit externe de la g tenseurs :

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)

Sur mon système, cela est environ sept fois plus rapide que la solution de Sven. Si le g ne change pas souvent, vous pouvez également mettre en cache le tenseur gggg tenseur. Si vous faites cela et que vous activez certaines micro-optimisations (inlining de la fonction tensordot code, pas de contrôles, pas de formes génériques), vous pouvez quand même le faire deux fois plus vite :

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))

Résultats de timeit sur mon ordinateur portable personnel (500 itérations) :

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

Les numéros sur ma machine de travail sont :

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

Answer 2

Voici comment le faire avec une seule boucle 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)

Certes, c'est un peu difficile à comprendre au premier abord, mais c'est un peu plus rapide :)

Answer 3

Grâce au travail acharné de M. Wiebe, la prochaine version de Numpy (qui sera probablement la 1.6) va rendre cela encore plus facile :

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

L'approche de Philipp est pour l'instant 3x plus rapide, mais il y a peut-être une marge d'amélioration. La différence de vitesse est probablement due au fait que tensordot est capable de dérouler toute l'opération en un seul produit matriciel qui peut être transmis à BLAS, et ainsi éviter une grande partie de l'overhead associé aux petits tableaux --- ce n'est pas possible pour la sommation générale d'Einstein, car toutes les opérations qui peuvent être exprimées sous cette forme ne se résolvent pas en un seul produit matriciel.

Answer 4

Par curiosité, j'ai comparé Cython mise en œuvre d'un code naïf de la question avec le code numpy de Réponse de @Philipp . Le code Cython est quatre fois plus rapide sur ma machine :

#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

Answer 5

J'ai pensé apporter un point de données relativement nouveau à ces repères, en utilisant perruche l'un des numpy -qui ont vu le jour au cours des derniers mois. (L'autre dont j'ai connaissance est numba mais je ne l'ai pas testé ici).

Après avoir traversé le processus d'installation quelque peu labyrinthique de LLVM, vous pouvez décorer de nombreux objets de pureté. numpy pour (souvent) accélérer leur exécution :

import numpy as np
import parakeet

@parakeet.jit
def rotT(T, g):
    # ...

J'ai seulement essayé d'appliquer le JIT au code d'Andrew dans la question originale, mais cela fonctionne plutôt bien (> 10x la vitesse) pour ne pas avoir à écrire de nouveau code :

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

Pour ces timings (sur mon ordinateur portable), j'ai exécuté chaque fonction dix fois, pour donner au JIT une chance d'identifier et d'optimiser les chemins de code chauds.

Il est intéressant de noter que les suggestions de Sven et Philipp sont toujours plus rapides de plusieurs ordres de grandeur !