Numpy einsum calcule le produit extérieur le long de l'axe

Question

J'ai deux tableaux numpy qui contiennent des matrices compatibles et je veux calculer le produit extérieur élément par élément de l'utilisation de numpy.einsum . Les formes des tableaux seraient:

A1 = (i,j,k)
A2 = (i,k,j) 

Par conséquent, les tableaux contiennent des imatrices de forme (k,j)et (j,k)respectivement.

Donc, donné A1contiendrait les matrices A,B,Cet A2contiendrait des matrices D,E,F, le résultat serait:

A3 = (A(x)D,B(x)E,C(x)F)

En (x)étant l'opérateur externe du produit.

Cela donnerait à ma compréhension basée sur cette réponse un tableau A3de la forme suivante:

A3 = (i,j*k,j*k)

Jusqu'à présent, j'ai essayé:

np.einsum("ijk, ilm -> ijklm", A1, A2)

Mais les formes résultantes ne correspondent pas correctement.

Pour vérifier la cohérence, je teste ceci:

A = np.asarray(([1,2],[3,4]))
B = np.asarray(([5,6],[7,8]))

AB_outer = np.outer(A,B)

A_vec = np.asarray((A,A))
B_vec = np.asarray((B,B))

# this line is not correct
AB_vec = np.einsum("ijk, ilm -> ijklm", A_vec,B_vec)

np.testing.assert_array_equal(AB_outer, AB_vec[0])

Cela génère actuellement une erreur d'assertion car ma notation einsum n'est pas correcte. Je suis également ouvert à toutes les suggestions qui peuvent résoudre ce problème et qui sont plus rapides ou aussi rapides que les nymphes einsum.

Answer 1

Nous pouvons prolonger les gradations et laisser broadcastingfaire le travail pour nous -

(A1[:,:,None,:,None]*A2[:,None,:,None,:]).swapaxes(2,3)

Exemple d'exécution -

In [46]: A1 = np.random.rand(3,4,4)
    ...: A2 = np.random.rand(3,4,4)

In [47]: out = (A1[:,:,None,:,None]*A2[:,None,:,None,:]).swapaxes(2,3)

In [48]: np.allclose(np.multiply.outer(A1[0],A2[0]), out[0])
Out[48]: True

In [49]: np.allclose(np.multiply.outer(A1[1],A2[1]), out[1])
Out[49]: True

In [50]: np.allclose(np.multiply.outer(A1[2],A2[2]), out[2])
Out[50]: True

L'équivalent avec np.einsumserait -

np.einsum('ijk,ilm->ijklm',A1,A2)

Answer 2

Vous pouvez calculer le résultat en cours d'exécution:

result = np.einsum('ijk,ikl->ijl', A1, A2)

J'ai vérifié le code ci-dessus sur les données de test suivantes:

A = np.arange(1, 13).reshape(3, -1)
B = np.arange(2, 14).reshape(3, -1)
C = np.arange(3, 15).reshape(3, -1)
D = np.arange(1, 13).reshape(4, -1)
E = np.arange(2, 14).reshape(4, -1)
F = np.arange(3, 15).reshape(4, -1)
A1 = np.array([A, B, C])
A2 = np.array([D, E, F])

Le résultat est:

array([[[ 70,  80,  90],
        [158, 184, 210],
        [246, 288, 330]],

       [[106, 120, 134],
        [210, 240, 270],
        [314, 360, 406]],

       [[150, 168, 186],
        [270, 304, 338],
        [390, 440, 490]]])

Calculez maintenant 3 "résultats partiels":

res_1 = A @ D
res_2 = B @ E
res_3 = C @ F

et vérifiez qu'ils sont exactement les mêmes que des sections consécutives du résultat.