Понимание einsum NumPy

схватывание идеи numpy.einsum() очень легко, если вы это интуитивно понимаете. Вот простое описание с матричным умножением в качестве примера.

использовать numpy.einsum(), вы должны пройти так называемый строка подстрочные в качестве аргумента, а затем на входных массивов.

предположим, у вас есть два 2D-массива,A и B, и вы хотите сделать умножение матрицы. Итак, вы:

np.einsum("ij, jk -> ik", A, B)

здесь индекс строки ij соответствует array A с индекс строки jk соответствует array B. Кроме того, самое главное отметить, что количество символов в каждом элементе индекс строки должны соответствуют размерам массива. (т. е. два символа для 2D-массивов, три символа для 3D-массивов и т. д.) И если вы повторите символы между индекс строки (j в нашем случае), то это означает, что вы хотите ein sum происходить вдоль этих измерений. Таким образом, они будут сокращены.

на индекс строки после этого ->, будет нашим результирующим массивом. Если вы оставьте его пустым, тогда все будет суммировано и скаляр будет возвращен в результате. То результирующий массив будет иметь размеры в соответствии с индекс строки. В нашем примере, это будет ik. Это интуитивно, потому что мы знаем, что для умножения матрицы количество столбцов в массиве A должно соответствовать количеству строк в массиве B что здесь происходит (т. е. мы кодируем этого знания по повторяя char j на индекс строки)

вот несколько примеров, иллюстрирующих использование np.einsum() в реализации некоторых общих тензора или НД-массив операции.

входы

In [197]: vec
Out[197]: array([0, 1, 2, 3])

In [198]: A
Out[198]: 
array([[11, 12, 13, 14],
       [21, 22, 23, 24],
       [31, 32, 33, 34],
       [41, 42, 43, 44]])


In [199]: B
Out[199]: 
array([[1, 1, 1, 1],
       [2, 2, 2, 2],
       [3, 3, 3, 3],
       [4, 4, 4, 4]])

1) умножение матриц (аналогично np.matmul(arr1, arr2))

In [200]: np.einsum("ij, jk -> ik", A, B)
Out[200]: 
array([[130, 130, 130, 130],
       [230, 230, 230, 230],
       [330, 330, 330, 330],
       [430, 430, 430, 430]])

2) извлекать элементы по главной диагонали (аналог np.diag(arr))

In [202]: np.einsum("ii -> i", A)
Out[202]: array([11, 22, 33, 44])

3) продукт Адамара (т. е. элементарное произведение двух массивов) (аналог arr1 * arr2)

In [203]: np.einsum("ij, ij -> ij", A, B)
Out[203]: 
array([[ 11,  12,  13,  14],
       [ 42,  44,  46,  48],
       [ 93,  96,  99, 102],
       [164, 168, 172, 176]])

4) элемент-мудрый квадратура (аналог np.square(arr) или arr ** 2)

In [210]: np.einsum("ij, ij -> ij", B, B)
Out[210]: 
array([[ 1,  1,  1,  1],
       [ 4,  4,  4,  4],
       [ 9,  9,  9,  9],
       [16, 16, 16, 16]])

5) Трассировка (т. е. сумма главных диагональных элементов) (аналог np.trace(arr))

In [217]: np.einsum("ii -> ", A)
Out[217]: 110

6) транспонировать матрицу (аналог np.transpose(arr))

In [221]: np.einsum("ij -> ji", A)
Out[221]: 
array([[11, 21, 31, 41],
       [12, 22, 32, 42],
       [13, 23, 33, 43],
       [14, 24, 34, 44]])

7) внешнее произведение (векторов) (аналог np.outer(vec1, vec2))

In [255]: np.einsum("i, j -> ij", vec, vec)
Out[255]: 
array([[0, 0, 0, 0],
       [0, 1, 2, 3],
       [0, 2, 4, 6],
       [0, 3, 6, 9]])

8) внутреннее произведение (векторов) (аналог np.inner(vec1, vec2))

In [256]: np.einsum("i, i -> ", vec, vec)
Out[256]: 14

9) сумма вдоль оси 0 (аналог np.sum(arr, axis=0))

In [260]: np.einsum("ij -> j", B)
Out[260]: array([10, 10, 10, 10])

10) сумма по оси 1 (аналог np.sum(arr, axis=1))

In [261]: np.einsum("ij -> i", B)
Out[261]: array([ 4,  8, 12, 16])

11) Умножение Матрицы Партии

In [287]: BM = np.stack((A, B), axis=0)

In [288]: BM
Out[288]: 
array([[[11, 12, 13, 14],
        [21, 22, 23, 24],
        [31, 32, 33, 34],
        [41, 42, 43, 44]],

       [[ 1,  1,  1,  1],
        [ 2,  2,  2,  2],
        [ 3,  3,  3,  3],
        [ 4,  4,  4,  4]]])

In [289]: BM.shape
Out[289]: (2, 4, 4)

# batch matrix multiply using einsum
In [292]: BMM = np.einsum("bij, bjk -> bik", BM, BM)

In [293]: BMM
Out[293]: 
array([[[1350, 1400, 1450, 1500],
        [2390, 2480, 2570, 2660],
        [3430, 3560, 3690, 3820],
        [4470, 4640, 4810, 4980]],

       [[  10,   10,   10,   10],
        [  20,   20,   20,   20],
        [  30,   30,   30,   30],
        [  40,   40,   40,   40]]])

In [294]: BMM.shape
Out[294]: (2, 4, 4)

12) Сумма по оси 2 (аналог np.sum(arr, axis=2))

In [330]: np.einsum("ijk -> ij", BM)
Out[330]: 
array([[ 50,  90, 130, 170],
       [  4,   8,  12,  16]])

13) суммировать все элементы в массиве (аналогично np.sum(arr))

In [335]: np.einsum("ijk -> ", BM)
Out[335]: 480

14) сумма по нескольким осям (т. е. маргинализация)
(похожие на np.sum(arr, axis=(axis0, axis1, axis2, axis3, axis4, axis6, axis7)))

# 8D array
In [354]: R = np.random.standard_normal((3,5,4,6,8,2,7,9))

# marginalize out axis 5 (i.e. "n" here)
In [363]: esum = np.einsum("ijklmnop -> n", R)

# marginalize out axis 5 (i.e. sum over rest of the axes)
In [364]: nsum = np.sum(R, axis=(0,1,2,3,4,6,7))

In [365]: np.allclose(esum, nsum)
Out[365]: True

15) Двойные Продукты Точки (аналог np.sum (hadamard-произведение) cf. 3)

In [772]: A
Out[772]: 
array([[1, 2, 3],
       [4, 2, 2],
       [2, 3, 4]])

In [773]: B
Out[773]: 
array([[1, 4, 7],
       [2, 5, 8],
       [3, 6, 9]])

In [774]: np.einsum("ij, ij -> ", A, B)
Out[774]: 124

16) 2D и 3D умножение массива

такой умножение может быть очень полезно при решении линейной системы уравнений ( Ax = b), где вы хотите, чтобы проверить результат.

# inputs
In [115]: A = np.random.rand(3,3)
In [116]: b = np.random.rand(3, 4, 5)

# solve for x
In [117]: x = np.linalg.solve(A, b.reshape(b.shape[0], -1)).reshape(b.shape)

# 2D and 3D array multiplication :)
In [118]: Ax = np.einsum('ij, jkl', A, x)

# indeed the same!
In [119]: np.allclose(Ax, b)
Out[119]: True

наоборот, если нужно использовать np.matmul() для этой проверки, мы должны сделать пару reshapes для достижения этого, как:

# reshape 3D array `x` to 2D, perform matmul
# then reshape the resultant array to 3D
In [123]: Ax_matmul = np.matmul(A, x.reshape(x.shape[0], -1)).reshape(x.shape)

# indeed correct!
In [124]: np.allclose(Ax, Ax_matmul)
Out[124]: True

бонус: подробнее по математике здесь:Эйнштейн-Суммирование и здесь: Тензорной Нотации

позволяет сделать 2 массива, с различными, но совместимыми размерами, чтобы выделить их взаимодействие

In [43]: A=np.arange(6).reshape(2,3)
Out[43]: 
array([[0, 1, 2],
       [3, 4, 5]])


In [44]: B=np.arange(12).reshape(3,4)
Out[44]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

ваш расчет принимает "точку" (сумму продуктов) a (2,3) с A (3,4) для получения массива (4,2). i является 1-м тусклым A, последний C; k последние B, 1-й C. j "потребляется" суммированием.

In [45]: C=np.einsum('ij,jk->ki',A,B)
Out[45]: 
array([[20, 56],
       [23, 68],
       [26, 80],
       [29, 92]])

это то же самое, что np.dot(A,B).T - это последний выход это транспонированный.

чтобы увидеть больше того, что происходит с j изменить C индексы ijk:

In [46]: np.einsum('ij,jk->ijk',A,B)
Out[46]: 
array([[[ 0,  0,  0,  0],
        [ 4,  5,  6,  7],
        [16, 18, 20, 22]],

       [[ 0,  3,  6,  9],
        [16, 20, 24, 28],
        [40, 45, 50, 55]]])

это можно также произвести с:

A[:,:,None]*B[None,:,:]

то есть добавить k измерение до конца A и i перед B, что приводит к массиву (2,3,4).

0 + 4 + 16 = 20, 9 + 28 + 55 = 92, etc; Sum on j и транспонировать, чтобы получить более ранний результат:

np.sum(A[:,:,None] * B[None,:,:], axis=1).T

# C[k,i] = sum(j) A[i,j (,k) ] * B[(i,)  j,k]

нашел NumPy: трюки торговли (Часть II) поучи

мы используем -> для указания порядка выходного массива. Поэтому подумайте о "ij, i - >j" как о левой стороне (LHS) и правой стороне (RHS). Любое повторение меток на LHS вычисляет элемент продукта мудро,а затем суммирует. Изменяя метку на стороне RHS (output), мы можем определить ось, по которой мы хотим продолжить относительно входного массива, т. е. суммирование вдоль оси 0, 1 и так далее.

import numpy as np

>>> a
array([[1, 1, 1],
       [2, 2, 2],
       [3, 3, 3]])
>>> b
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> d = np.einsum('ij, jk->ki', a, b)

для того, чтобы рассчитать продукт и выровнять j оси надо добавить ось a. (b будет транслироваться по(?) первая ось)

a[i, j, k]
   b[j, k]

>>> c = a[:,:,np.newaxis] * b
>>> c
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 0,  2,  4],
        [ 6,  8, 10],
        [12, 14, 16]],

       [[ 0,  3,  6],
        [ 9, 12, 15],
        [18, 21, 24]]])

j отсутствует с правой стороны, поэтому мы суммируем j что вторая ось массива 3x3x3

>>> c = c.sum(1)
>>> c
array([[ 9, 12, 15],
       [18, 24, 30],
       [27, 36, 45]])

наконец, индексы (в алфавитном порядке) перевернуты с правой стороны, поэтому мы транспонируем.

>>> c.T
array([[ 9, 18, 27],
       [12, 24, 36],
       [15, 30, 45]])

>>> np.einsum('ij, jk->ki', a, b)
array([[ 9, 18, 27],
       [12, 24, 36],
       [15, 30, 45]])
>>>