numpy.einsum

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

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

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

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

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

Параметры

subscriptsstr

Указывает индексы для суммирования в виде списка меток индексов, разделенных запятыми. Неявное (классическое эйнштейновское суммирование) вычисление выполняется, если не указан явный индикатор «->», а также метки индексов точного выходного формата.

operandslist of array_like

Это массивы для операции.

outndarray, 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.

Возвращает

outputndarray

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

См. также

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 по умолчанию не допускает распрoстранения. Кроме того np.einsum('ij,jh->ih', a, b) напрямую определяет порядок меток индексов выходного значения и, следовательно, возвращает результат матричного умножения, в отличие от примера выше в неявном режиме.

Для включения и управления распрoстранением используйте эллипсис. По умолчанию распрoстранение в стиле 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)

Базовая einsum: ~1520мс (измерялось на 3,1ГГц Intel i5.)

>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a)

Неоптимальная einsum (из-за повторного вычисления времени пути): ~330мс

>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='optimal')

Жадный einsum (более быстрое приближение оптимального пути): ~160мс

>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='greedy')

Оптимальная einsum (лучший способ использования в некоторых случаях): ~110мс

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