Numpy einsum вычисляет внешний продукт вдоль оси

Question

У меня есть два массива numpy, которые содержат совместимые матрицы, и я хочу вычислить поэлементный внешний продукт использования numpy.einsum . Формы массивов будут:

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

Поэтому массивы содержат iматрицы формы (k,j)и (j,k)соответственно.

Таким образом, данный A1будет содержать матрицы A,B,Cи A2содержать матрицы D,E,F, результат будет:

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

С (x)того внешний оператор продукта.

На основании этого ответа я понял бы, что это массив A3следующей формы:

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

Пока я пробовал:

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

Но получившиеся формы не подходят правильно.

В качестве проверки работоспособности я тестирую это:

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

В настоящее время это вызывает ошибку утверждения, поскольку моя нотация einsum неверна. Я также открыт для любых предложений, которые могут решить эту проблему и будут быстрее или столь же быстрыми, как nymphs einsum.

Answer 1

Мы можем увеличить яркость и позволить broadcastingсделать всю работу за нас -

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

Пробный прогон -

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

Эквивалент с np.einsumбудет -

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

Answer 2

Вы можете вычислить результат, запустив:

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

Я проверил приведенный выше код на следующих тестовых данных:

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

Результат:

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

Теперь вычислим 3 «частичных результата»:

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

и убедитесь, что они совпадают с последовательными частями результата.