Einsum

L’autre jour, j’ai rencontré Einstein se promenant incognito dans un bout de code.

J’ai commencé par ne rien comprendre à une ligne de Python trouvée dans un algo de traitement d’image:

norms = np.einsum(‘ij,ij->i’, X, X)

Alors je suis allé voir la doc de la fonction einsum de NumPy à laquelle je n’ai rien compris non plus, sauf que le bout du nez d’Albert commençait à apparaître via une mystérieuse « convention de sommation d’Einstein » qui serait « un raccourci de notation utile pour la manipulation des équations concernant des coordonnées ».

Par exemple, une simple multiplication matricielle de deux matrices A et B s’écrit   $$A^i{}_j \times B^j{}_k = C^i{}_k$$ avec la notation d’Einstein, et C=np.einsum(‘ij,jk->ik’, A, B) avec la fonction Python correspondante.

On voit que la chaîne de caractères ‘ij,jk->ik’ correspond aux indices des lignes et colonnes des matrices qui figurent en indice et « exposant » dans la notation d’Einstein, dans le même ordre que dans la représentation de la multiplication matricielle ci-dessous:

Pour une simple multiplication matricielle on ne voit pas trop l’avantage sur la notation $$A \times B = C$$ , mais ajoutons les 3 règles supplémentaires de la notation d’Einstein appliqués à la fonction einsum (traduit de [2])

1. Les indices qui se répètent entre tableaux (qui peuvent avoir N dimensions) signifient que les valeurs selon ces axes doivent être multipliées. Les produits font les valeurs du tableau de sortie.
   Dans le cas ci-dessus on a utilisé l’indice ‘j’ une fois pour A et une fois pour B, ce qui signifie qu’on multiplie chaque ligne de A avec chaque colonne de B
2. Une lettre omise dans le résultat signifie que les valeurs le long de cette axe doivent être sommées.
   Dans le produit matriciel ci-dessus, ‘j’ ne figure pas dans les indices du tableau de sortie. Cette omission somme le long de cette « troisième direction en réduisant de 1 la dimension du tableau final. Si la signature avait été ‘ij,jk->ijk’, on aurait obtenu une matrice 3x3x3 des produits, et si on n’avait écrit aucun indice mais gardé la flèche ‘ij,jk->’, on aurait obtenu la somme de tous les éléments de C
3. On peut obtenir les axes non sommés dans n’importe quel ordre.
   Si on omet la flèche, einsum arrange les indices qui n’apparaissent qu’une fois dans l’ordre alphabétique, donc on aurait pu écrire juste ‘ij,jk’ au lieu de ‘ij,jk->ik’ pour le produit de matrices « normal » ci-dessus.
   Mais on aurait pu écrire ‘ij,jk->ki’ si on avait souhaité obtenir la transposée du résultat.

Avec ces astuces, beaucoup (toutes?) les opérations vectorielles, matricielles et sur des tableaux à N dimensions peuvent être effectuées par einsum

Les mathématiciens apprécieront l’élégance de la notation, les physiciens pourront manipuler le tenseur de Ricci et d’autres objets mathématiques avec plus de facilité, mais en quoi la fonction einsum est-elle utile aux informaticiens ?

D’abord, elle permet d’effectuer en une seule fois des opérations qui auraient nécessité de stocker des matrices intermédiaires sans ça. De plus le code de la fonction einsum peut être fortement vectorisé puisqu’il effectue essentiellement des produits puis des sommes sur des vecteurs. Ce n’est en définitive que l’indexation des tableaux qui change en fonction de l’opération désirée, et le tout peut être optimisé dans certains cas en jouant avec la représentation interne des matrices et la commutativité des opérations. Grâce à ceci, la plupart des opérations listées dans le tableau ci-dessus sont entre 1 et 4 fois plus rapides avec einsum qu’avec leur équivalent Numpy, jamais plus lentes. Et pour des opérations plus complexes comme le calcul de la matrice de Fock, l’accélération peut être d’un facteur 10 ou plus [4].

Si vous n’avez pas tout suivi, retenez juste ceci:

1. Einstein n’était pas nul en maths du tout. Il était un surdoué en maths qui a préféré faire de la physique, mais le fait que sa « convention de sommation » apparaisse dans les langages de programmation un siècle plus tard confirme le génie du bonhomme
2. Que Python est LE langage de programmation scientifique qui s’impose. Oubliez Matlab, qui n’offre même pas de fonction einsum …

Références

1. [altmetric doi= »10.1016/S0898-1221(02)00210-9″ float= »right »] Åhlander, K. (2002). « Einstein summation for multidimensional arrays », Computers and Mathematics with Applications, 44, 1007–1017.  doi:10.1016/S0898-1221(02)00210-9 [pdf]
2. Alex Riley, « A basic introduction to NumPy’s einsum« 
3. Einstein summation sur MathWorld
4. Einsum tutorial sur GitHub
5. Code source en C de la fonction einsum de NumPy (pas simple…)