numpy.einsum

numpy.einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', optimize=False) [source]

Вычисляет соглашение об эйнштейновской сумме по операндам.

Используя соглашение об эйнштейновской сумме, многие распространенные многомерные линейно-алгебраические операции над массивами могут быть представлены в простом виде. В неявном режиме einsum вычисляет эти значения.

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

См. примечания и примеры для уточнения.

Параметры:

subscripts : str Указывает подписи для суммирования как список подписей, разделенных запятыми. Выполняется неявное (классическое эйнштейновское суммирование) вычисление, если не указан явный индикатор ‘->’, а также метки подписей точной формы вывода. operands : list of array_like Это массивы для операции. out : ndarray, optional Если указано, вычисление выполняется в этом массиве. dtype : {data-type, None}, optional Если указано, принудительно использует указанный тип данных для вычисления. Обратите внимание, что вам может потребоваться также указать более либеральный параметр casting для разрешения преобразований. По умолчанию None. order : {‘C’, ‘F’, ‘A’, ‘K’}, optional Управляет расположением памяти вывода. ‘C’ означает, что оно должно быть непрерывным в стиле C. ‘F’ означает, что оно должно быть непрерывным в стиле Fortran, ‘A’ означает, что оно должно быть ‘F’, если все входные данные являются ‘F’, иначе ‘C’. ‘K’ означает, что оно должно быть как можно ближе к расположению входных данных, включая произвольную перестановку осей. По умолчанию ‘K’. casting : {‘no’, ‘equiv’, ‘safe’, ‘same_kind’, ‘unsafe’}, optional Управляет тем, какие преобразования типов данных могут произойти. Не рекомендуется устанавливать это значение в ‘unsafe’, так как это может негативно повлиять на накопления. ‘no’ означает, что типы данных вообще не должны преобразовываться. ‘equiv’ означает, что разрешены только изменения порядка байтов. ‘safe’ означает, что разрешены только преобразования, которые могут сохранить значения. ‘same_kind’ означает, что разрешены только безопасные преобразования или преобразования в рамках типа, например, от float64 к float32. ‘unsafe’ означает, что любые преобразования данных могут быть выполнены. По умолчанию ‘safe’. optimize : {False, True, ‘greedy’, ‘optimal’}, optional Управляет тем, должна ли происходить промежуточная оптимизация. Если False, оптимизация не будет происходить, а True по умолчанию будет использовать алгоритм ‘greedy’. Также принимает явный список сокращений из функции np.einsum_path. Подробнее см. np.einsum_path. По умолчанию False.

Возвращает:

output : ndarray Вычисление на основе соглашения об эйнштейновской сумме.

См. также

einsum_path, dot, inner, outer, tensordot, linalg.multi_dot

Примечания

Новое в версии 1.6.0.

Соглашение об эйнштейновской сумме может использоваться для вычисления многих многомерных линейно-алгебраических операций над массивами. einsum предоставляет краткий способ представления этих операций.

Ниже приводится неполный список этих операций, которые могут быть вычислены с помощью einsum, вместе с примерами:

• След массива, numpy.trace.
• Возвращение диагонали, numpy.diag.
• Суммирование по осям массива, numpy.sum.
• Транспонирование и перестановки, numpy.transpose.
• Умножение матриц и скалярное произведение, numpy.matmul numpy.dot.
• Векторное внутреннее и внешнее произведение, numpy.inner numpy.outer.
• Вещание, поэлементное и скалярное умножение, numpy.multiply.
• Сжатие тензоров, numpy.tensordot.
• Цепочечные операции над массивами в эффективном порядке вычислений, numpy.einsum_path.

Строка подписей — это список подписей, разделенных запятыми, где каждая подпись относится к измерению соответствующего операнда. Когда подпись повторяется, она суммируется, поэтому np.einsum('i,i', a, b) эквивалентно np.inner(a,b). Если подпись встречается только один раз, она не суммируется, поэтому np.einsum('i', a) создает представление a без изменений. Дополнительный пример np.einsum('ij,jk', a, b) описывает традиционное умножение матриц и эквивалентен np.matmul(a,b). Повторяющиеся подписи в одном операнде берут диагональ. Например, np.einsum('ii', a) эквивалентно np.trace(a).

В неявном режиме выбранные подписи важны, поскольку оси вывода упорядочиваются в алфавитном порядке. Это означает, что np.einsum('ij', a) не влияет на 2D массив, а np.einsum('ji', a) транспонирует его. Кроме того, np.einsum('ij,jk', a, b) возвращает умножение матриц, тогда как np.einsum('ij,jh', a, b) возвращает транспонированное умножение, так как подпись ‘h’ предшествует подписи ‘i’.

В явном режиме вывод может быть непосредственно контролируется путем указания подписей вывода. Это требует идентификатора ‘->’, а также списка меток подписей вывода. Эта функция повышает гибкость функции, поскольку суммирование можно отключить или принудить, если это необходимо. Вызов np.einsum('i->', a) похож на np.sum(a, axis=-1), а np.einsum('ii->i', a) похож на np.diag(a). Разница в том, что einsum по умолчанию не допускает вещание. Кроме того, np.einsum('ij,jh->ih', a, b) напрямую определяет порядок меток подписей вывода и, следовательно, возвращает умножение матриц, в отличие от примера выше в неявном режиме.

Для включения и управления вещанием используйте эллипсис. Вещание в стиле NumPy по умолчанию выполняется путем добавления эллипса слева от каждого члена, например, np.einsum('...ii->...i', a). Для вычисления следа по первой и последней осям можно выполнить np.einsum('i...i', a), или для выполнения умножения матриц-матриц с левыми индексами вместо правых можно выполнить np.einsum('ij...,jk...->ik...', a, b).

Когда имеется только один операнд, не суммируются никакие оси, и не предоставлен параметр вывода, вместо нового массива возвращается представление операнда. Таким образом, взятие диагонали как np.einsum('ii->i', a) создаёт представление (изменено в версии 1.10.0).

einsum также предоставляет альтернативный способ предоставления подписей и операндов в виде einsum(op0, sublist0, op1, sublist1, ..., [sublistout]). Если форма вывода не указана в этом формате einsum будет рассчитана в неявном режиме, в противном случае - явно. Примеры ниже имеют соответствующие вызовы einsum с методами с двумя параметрами.

Новое в версии 1.10.0.

Представления, возвращаемые функцией einsum, теперь доступны для записи, когда входной массив доступен для записи. Например, np.einsum('ijk...->kji...', a) теперь будет иметь такой же эффект, как np.swapaxes(a, 0, 2), и np.einsum('ii->i', a) вернет доступное для записи представление диагонали 2D массива.

Новое в версии 1.12.0.

Добавлен аргумент optimize, который оптимизирует порядок сжатия выражения einsum. Для сжатия с тремя или более операндами это может значительно повысить вычислительную эффективность за счёт большего объёма памяти во время вычисления.

Обычно применяется алгоритм «жадный», который, как показали эмпирические тесты, в большинстве случаев возвращает оптимальный путь. В некоторых случаях «оптимальный» вернёт превосходный путь путём более дорогостоящего, исчерпывающего поиска. Для итеративных вычислений может быть целесообразно вычислить оптимальный путь один раз и использовать этот путь повторно, предоставив его в качестве аргумента. Пример показан ниже.

См. numpy.einsum_path для более подробной информации.

Примеры

>>> a = np.arange(25).reshape(5,5)
>>> b = np.arange(5)
>>> c = np.arange(6).reshape(2,3)

След матрицы:

>>> np.einsum('ii', a)
60
>>> np.einsum(a, [0,0])
60
>>> np.trace(a)
60

Извлечь диагональ (требуется явная форма):

>>> np.einsum('ii->i', a)
array([ 0,  6, 12, 18, 24])
>>> np.einsum(a, [0,0], [0])
array([ 0,  6, 12, 18, 24])
>>> np.diag(a)
array([ 0,  6, 12, 18, 24])

Суммирование по оси (требуется явная форма):

>>> np.einsum('ij->i', a)
array([ 10,  35,  60,  85, 110])
>>> np.einsum(a, [0,1], [0])
array([ 10,  35,  60,  85, 110])
>>> np.sum(a, axis=1)
array([ 10,  35,  60,  85, 110])

Для многомерных массивов суммирование по одной оси может быть выполнено с использованием эллипсиса:

>>> np.einsum('...j->...', a)
array([ 10,  35,  60,  85, 110])
>>> np.einsum(a, [Ellipsis,1], [Ellipsis])
array([ 10,  35,  60,  85, 110])

Вычислить транспонированную матрицу или переупорядочить любое количество осей:

>>> np.einsum('ji', c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.einsum('ij->ji', c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.einsum(c, [1,0])
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.transpose(c)
array([[0, 3],
       [1, 4],
       [2, 5]])

Внутренние произведения векторов:

>>> np.einsum('i,i', b, b)
30
>>> np.einsum(b, [0], b, [0])
30
>>> np.inner(b,b)
30

Умножение матрицы на вектор:

>>> np.einsum('ij,j', a, b)
array([ 30,  80, 130, 180, 230])
>>> np.einsum(a, [0,1], b, [1])
array([ 30,  80, 130, 180, 230])
>>> np.dot(a, b)
array([ 30,  80, 130, 180, 230])
>>> np.einsum('...j,j', a, b)
array([ 30,  80, 130, 180, 230])

Расширение и умножение на скаляр:

>>> np.einsum('..., ...', 3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.einsum(',ij', 3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.einsum(3, [Ellipsis], c, [Ellipsis])
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.multiply(3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])

Внешнее произведение векторов:

>>> np.einsum('i,j', np.arange(2)+1, b)
array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])
>>> np.einsum(np.arange(2)+1, [0], b, [1])
array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])
>>> np.outer(np.arange(2)+1, b)
array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])

Сжатие тензора:

>>> a = np.arange(60.).reshape(3,4,5)
>>> b = np.arange(24.).reshape(4,3,2)
>>> np.einsum('ijk,jil->kl', a, b)
array([[ 4400.,  4730.],
       [ 4532.,  4874.],
       [ 4664.,  5018.],
       [ 4796.,  5162.],
       [ 4928.,  5306.]])
>>> np.einsum(a, [0,1,2], b, [1,0,3], [2,3])
array([[ 4400.,  4730.],
       [ 4532.,  4874.],
       [ 4664.,  5018.],
       [ 4796.,  5162.],
       [ 4928.,  5306.]])
>>> np.tensordot(a,b, axes=([1,0],[0,1]))
array([[ 4400.,  4730.],
       [ 4532.,  4874.],
       [ 4664.,  5018.],
       [ 4796.,  5162.],
       [ 4928.,  5306.]])

Возвращаемые массивы с возможностью записи (с версии 1.10.0):

>>> a = np.zeros((3, 3))
>>> np.einsum('ii->i', a)[:] = 1
>>> a
array([[ 1.,  0.,  0.],
       [ 0.,  1.,  0.],
       [ 0.,  0.,  1.]])

Пример использования эллипса:

>>> a = np.arange(6).reshape((3,2))
>>> b = np.arange(12).reshape((4,3))
>>> np.einsum('ki,jk->ij', a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])
>>> np.einsum('ki,...k->i...', a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])
>>> np.einsum('k...,jk', a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])

Цепочечные операции с массивами. Для более сложных сжатий ускорение может быть достигнуто, вычисляя последовательно «жадный» путь или предварительно вычисляя «оптимальный» путь и повторно применяя его, используя вставку einsum_path (с версии 1.12.0). Улучшения производительности могут быть особенно значительными с большими массивами:

>>> a = np.ones(64).reshape(2,4,8)
# Basic `einsum`: ~1520ms  (benchmarked on 3.1GHz Intel i5.)
>>> for iteration in range(500):
...     np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a)
# Sub-optimal `einsum` (due to repeated path calculation time): ~330ms
>>> for iteration in range(500):
...     np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='optimal')
# Greedy `einsum` (faster optimal path approximation): ~160ms
>>> for iteration in range(500):
...     np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='greedy')
# Optimal `einsum` (best usage pattern in some use cases): ~110ms
>>> path = np.einsum_path('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='optimal')[0]
>>> for iteration in range(500):
...     np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize=path)