Rotation tensorielle rapide avec NumPy

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

Certes, c'est un peu difficile à comprendre au premier coup d'œil, mais c'est un peu plus rapide :)

grâce au travail acharné de M. Wiebe, la prochaine version de Numpy (qui sera probablement 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 actuellement 3 fois plus rapide, mais il y a peut-être place à amélioration. La différence de vitesse est probablement dû principalement à tensordot étant en mesure de dérouler l'ensemble de l'opération comme un produit de matrice simple qui peut être transmise à BLAS, et donc en évitant une grande partie de la charge associée à de petites tableaux - - - - ceci 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 de matrice.

par curiosité j'ai comparé Cython mise en œuvre d'un code naïf de la question avec le code nul 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

j'ai pensé que j'apporterais un point de données relativement nouveau à ces benchmarks, en utilisant parakeet , l'un des numpy - compilateurs JIT conscients qui a vu le jour dans les derniers mois. (L'autre que je connais est numba , mais je ne l'ai pas testé ici.)

après que vous le faites à travers le processus d'installation quelque peu labyrinthine pour LLVM, vous pouvez décorer de nombreuses pures - numpy fonctions pour (souvent) accélérer leur performance:

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 il fait assez bien (>10x speedup) pour n'avoir à écrire aucun nouveau code que ce soit:

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 chronométrages (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.

fait intéressant, les suggestions de Sven et Philipp sont encore des ordres de grandeur plus vite !

Approche Prospective et solution de code

pour l'efficacité de la mémoire et par la suite l'efficacité de la performance, nous pourrions utiliser la multiplication par étapes de la matrice tensorielle.

pour illustrer les étapes impliquées, utilisons la plus simple des solutions avec np.einsum de @pv. -

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

comme vu, nous perdons la première dimension de g avec tenseur-multiplication entre ses quatre variantes et T .

faisons ces réductions de somme pour les multiplications de matrice tensorielle par étapes. Commençons par la première variante de g et T :

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

ainsi, nous nous retrouvons avec un tenseur de dimensions en notation chaîne: bcdi . Les étapes suivantes consisteraient à réduire ce tenseur par rapport au reste des trois variantes g utilisées dans l'original einsum de mise en oeuvre. Par conséquent, la réduction suivante serait de -

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

comme vu, nous avons perdu les deux premières dimensions avec les notations de chaîne: a , b . Nous le continuons pendant deux autres étapes pour se débarrasser de c et d aussi et serait laissé avec ijkl comme la sortie finale, comme so -

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

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

maintenant, nous pourrions utiliser np.tensordot pour ces réductions de somme, qui serait beaucoup plus efficace.

mise en œuvre finale

ainsi, portant sur np.tensordot , nous aurions la mise en œuvre finale comme si -

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

essai D'exécution

nous allons tester toutes les approches basées sur NumPy affiché à travers d'autres postes pour résoudre le problème sur la performance.

Approches comme des fonctions -

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

Timings avec l'original de la taille des jeux de données -

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

Comme il est mentionné au début de ce post, nous essayons d'atteindre l'efficacité de mémoire et donc de meilleures performances. Testons cela en augmentant la taille des ensembles de données -

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