Соглашение Эйнштейна и einsum

Удивительное дело, но в русскоязычном сегменте интернета почти нет материала, разъясняющего понятным языком соглашение Эйнштейна о суммировании. Не менее удивительно то, что материалов, позволяющих понять принцип работы функции einsum в русскоязычном интернете ещё меньше. На английском есть довольно развёрнутый ответ о работе einsum на stack overflow, а на русском только некоторое число сайтов, предоставляющих кривой перевод этого самого ответа. Хочу исправить эту проблему с недостатком материалов, и всех, кому интересно приглашаю к прочтению!


Обсуждаем соглашение Эйнштейна

Прежде всего отмечу, что соглашение Эйнштейна чаще всего используется в тензорном анализе и его приложениях, поэтому дальше в статье будет несколько референсов к тензорам.
Когда вы только начинаете работать с тензорами, вас может смутить, что кроме привычных подстрочных индексов, используются также и надстрочные индексы, которые по началу вообще можно принять за возведение в степень. Пример:
"а с верхним индексом i" будет записано как a^i, а "a в квадрате с верхним индексом i" будет записываться (a^i)^2. Возможно, по-началу это вводит в заблуждение и кажется неудобным, но со временем можно привыкнуть.

Соглашение: далее в статье объекты вида a_ix_iилиa_ix^i я буду называть термами.

О чём вообще соглашение Эйнштейна?

Соглашение Эйнштейна призвано уменьшить число знаков суммирования в выражении. Есть три простых правила, определяющие, насколько то или иное выражение корректно записано в нотации Эйнштейна.

Правило № 1: Суммирование ведётся по всем индексам, повторяющимся дважды в одном терме.

Пример: рассмотрим выражение следующего вида:

\sum_{i = 1}^3 a_ix_i = a_1x_1 + a_2x_2 + a_3x_3

С использованием соглашения Эйнштейна это выражение может быть переписано так:

a_ix_i \text{ или } a_ix^i

Таким образом мы избавляемся от знака суммы, и просто пишем единственный терм. Обратим внимание, что в этом терме индекс i повторяется дважды, а значит, в соответствие с первым правилом мы понимаем, что суммирование ведётся по индексу i, а точнее, по всем возможным значениям, которые принимает этот индекс.


Рассмотрим ещё один пример: пусть нам нужно умножить матрицу A \in \mathbb{R}^{m\times n} на вектор v \in \mathbb{R}^{n}. Результатом будет являться вектор b  \in \mathbb{R}^{m}. По определению:

b_i = \sum\limits_{j = 1}^n A_{ij} v_j, ~ i = 1, \ldots, m

Соглашение Эйнштейна позволяет избавиться от знака суммы:

b_i = A_{ij}v_{j} = A_{ij}v^{j}

Заметим, что в терм индекс i входит один раз, а индекс j входит два раза, а значит, суммирование ведётся по индексу j.

Определение 1. Индекс, который входит в терм дважды, называется фиктивным индексом.

Определение 2. Свободным индексом назовём все индексы в терме, не являющие фиктивными.

Отметим, что каждый фиктивный индекс может быть заменён любым другим фиктивным индексом, при условии, что

  1. Новый фиктивный индекс не входит в множество свободных индексов терма.

  2. Новый фиктивный индекс принимает то же множество значений, что и старый фиктивный индекс.

Чтобы объяснить проще, рассмотрим следующий код на языке Python:

for i in range(M):
    for j in range(N):
        b[i] += A[i, j] * v[j]

Этот код кратко описывает процесс умножения матрицы на вектор, а точнее, этот пример. Здесь индекс j является фиктивным, а индекс i – свободным. Суммирование в соглашении Эйнштейна ведётся по фиктивным индексам. Имя переменной j мы можем заменить на любое другое.

Правило № 2. В каждом терме не может встречаться более двух одинаковых индексов.

Второе правило говорит нам, что мы можем написать a_{ij}b_{ij}, но не можем написать a_{ii}b_{ij} или a_{ij}b_{jj}, несмотря на то, что на практике такие выражения всё же имеют смысл.
Больше примеров:

a_i^i– здесь i является фиктивным индексом, т.к. повторяется дважды;

a_i^{jj}– здесь i является свободным индексом, а j – фиктивным;

a_{ii}^{jj}– здесь и i, и j являются фиктивными индексами;

a_{ij}^{ij}– здесь и i, и j являются фиктивными индексами;

a_{ii}^{ij}– не правильно по второму правилу (индекс i входит в терм трижды);

Из примеров выше можно заключить, что когда мы считаем число вхождений индексов в терм, мы не делаем разницы между верхними и нижними индексами, и считаем их вместе. Ещё один важный пример: когда мы видим выражение следующего вида

a_{ij}b_{i} + a_{ji}b_{j}

Мы должны понимать, что это выражение записано верно, и не противоречит второму правилу. Действительно, если посчитать все вхождения индексов, то получится, что индекс i входит 3 раза, как и индекс j, но в выражении записано два терма, а не один, и если посчитать вхождение индексов в каждый терм отдельно (как того и требует второе правило), то мы увидим, что ничего не нарушается.

Правило № 3. В уравнениях, записанных с использованием соглашения Эйнштейна свободные индексы слева и свободные индексы справа должны совпадать.

Рассмотрим несколько примеров для закрепления этого правила:

b_i = A_{ij}v_{j}– этот пример мы уже рассматривали выше, здесь i является свободным индексом левой части уравнения, и свободным индексом правой части уравнения;

a_i = A_{ki}B_{kj}x_{j} + C_{ik}u_{k}– пример посложнее. Посчитаем вхождения индексов для каждого терма: в первый терм правой части k и j входят дважды, значит, они являются фиктивными индексами, i входит один раз, значит, является свободным. Во второй терм правой части k входит два раза, i – один, значит, k – фиктивный, i – свободный. В левой части индекс i входит один раз, а значит, является свободным. Итог: индекс i является свободным для обеих частей уравнения, а значит, правило 3 выполнено.

Рассмотрим так же несколько примеров, в которых третье правило не выполняется:

x_i = A_{ij}– слева i является свободным индексом, но справа свободны индексы i и j;

x_j=A_{ik}u_k– слева свободен индекс j, но справа свободен индекс i. Свободные индексы не совпадают;

x_i = A_{ik}u_k + c_j– здесь слева свободен индекс i, а справа свободны индексы i, j;

Пример упрощения сложного выражения с помощью соглашения Эйнштейна: тензорный поезд

Пусть A – пятимерный тензор. Тогда утверждается, что он может быть представлен в следующем виде:

A_{i_1i_2i_3i_4i_5} = \sum_{j_4=1}^{R_4}\sum_{j_3=1}^{R_3}\sum_{j_2=1}^{R_2}\sum_{j_1=1}^{R_1}G^{(1)}_{i_1j_1}G^{(2)}_{j_1i_2j_2}G^{(3)}_{j_2i_3j_3}G^{(4)}_{j_3i_4j_4}G^{(5)}_{j_4i_5}

Там сейчас не очень важно, что из себя представляется каждая G^{(k)}, и что такое R_i. Наша задача сейчас – исключительно синтаксическая игра. Нужно упростить выражение, особо не вникая в смысл происходящего.
Прежде всего видно, что свободными индексами являются i_1, i_2, i_3, i_4, i_5, а фиктивными, соответственно индексы j_1, j_2, j_3, j_4. Расположим индексы в соседних множителях так, чтобы в первом множителе индекс, по которому идёт суммирование, стоял снизу, а во втором тот же самый индекс стоял сверху. Так же заметим, что множителями являются тензоры G^{(k)}, и у них в верхнем регистре уже стоит (k). Чтобы повысить читаемость, будем оборачивать множители в скобки, и только потом ставить индексы. Само же упрощённое выражение переписывается из исходного почти дословно:

A_{i_1i_2i_3i_4i_5}=\left(G^{(1)}\right)_{i_1j_1}\left(G^{(2)}\right)_{i_2j_2}^{j_1}\left(G^{(3)}\right)_{i_3j_3}^{j_2}\left(G^{(4)}\right)_{i_4j_4}^{j_3}\left(G^{(5)}\right)_{i_5}^{j_4}

Ура, мы научились упрощать сложные выражения с помощью соглашения Эйнштейна!

Обсуждаем einsum

einsum это функция, присутствующая в нескольких популярных библиотеках для Python (NumPy, TensorFlow, PyTorch). Во всех библиотеках, в которых эта функция реализована, она работает одинаково (с точностью до функционала структур, определённых в конкретной библиотеке), поэтому нет смысла рассматривать один и тот же пример в разных библиотеках, достаточно рассказать про einsum в одной конкретной библиотеке. Далее в статье я буду использовать NumPy. einsum применяет соглашение Эйнштейна о суммировании к переданным массивам. Функция принимает множество опциональных аргументов, про них лучше почитать в документации, мы же сейчас разберём, как передавать шаблон, по которому функция будет применять соглашение Эйнштейна.

Рассмотрим сразу такой пример: пусть A \in \mathbb{R}^{3\times5}, B \in \mathbb{R}^{5\times2} – две матрицы, и мы хотим их перемножить. Результатом будет матрица M \in \mathbb{R}^{3\times2}, которую мы можем записать следующим образом, используя определение матричного умножения и соглашение Эйнштейна:

M_{ij}=\sum_{k=1}^{5}A_{ik}B_{kj} = A_{ik}B_{kj}

Теперь пусть мы хотим перемножить их программно. Ну, это можно довольно просто сделать с помощью трёх вложенных циклов:

M = np.zeros((3, 2))
for i in range(3):
    for j in range(2):
        for k in range(5):
            M[i, j] += A[i, k] * B[k, j]

Либо, используя функцию einsum можно написать это произведение в одну строчку:

M = np.einsum("ik,kj->ij", A, B)

Разберёмся, что за магия происходит в этой строчке. einsum принимает один обязательный аргумент: шаблон, по которому будет применено соглашение Эйнштейна. Шаблон этот выглядит так:

"{индексы, определяющие размерность первого массива},{индексы, определяющие размерность второго массива}->{индексы, определяющие размерность результирующего массива}"

Поведение шаблона einsum определяется следующими правилами:

  • Если один и тот же индекс встречается слева и справа от запятой (до стрелочки), то суммирование будет вестись по этому индексу;

  • Если после стрелочки ничего не написано, то суммирование произойдёт по всем встреченным осям;

  • Никакой индекс не должен встречаться 3 и более раз;

Таким образом мы видим, что einsum очень естественно поддерживает понятие свободных и фиктивных индексов, а также первые два правила, которые мы вводили, пока обсуждали соглашение Эйнштейна. Кроме того, как выражение, написанное с помощью соглашения Эйнштейна, может быть развёрнуто с помощью введения знаков суммы, так и функция einsum может быть развёрнута с помощью нескольких вложенных циклов. Это может быть очень удобно на первых порах, пока не сформируется устойчивое понимание einsum.

Рассмотрим теперь некоторое количество примеров разной степени сложности, чтобы закрепить понимание einsum:

Одна einsum, чтобы править всеми

Пример 1. Сумма всех значений вектора:

vector = np.array([1, 2, 3, 4, 5])
result = np.einsum("i->", vector)
print(result)
Output

15

Пример 2. Сумма всех значений матрицы:

matrix = np.array([[1, 2], [3, 4], [5, 6]])
result = np.einsum("ij->", matrix)
print(result)
Output

21

Пример 3. Сумма значений по столбцам:

matrix = np.array([[1, 2], [3, 4], [5, 6]])
result = np.einsum("ij->j", matrix)
print(result)
Output

[9, 12]

Пример 4. Сумма значений по строкам:

matrix = np.array([[1, 2], [3, 4], [5, 6]])
result = np.einsum("ij->i", matrix)
print(result)
Output

[3, 7, 11]

Пример 5. Транспонирование (я об этом не написал, но оси, по которым суммирования не произошло, мы можем возвращать в любом порядке):

matrix = np.array([[1, 2], [3, 4], [5, 6]])
result = np.einsum("ij->ji", matrix)
print(result)
Output

[[1, 3, 5], [2, 4, 6]]

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

matrix = np.array([[1, 2], [3, 4], [5, 6]])
vector = np.array([[1, 2]])
result = np.einsum("ij,kj->ik", matrix, vector)
print(result)

Заметим, что вектор имеет форму 1 \times 2, и чтобы умножить матрицу на него по правилам, его нужно было бы транспонировать. Однако с помощью einsum мы можем задать ось, по которой будет вестись суммирование, и немного выиграть по памяти, не создавая копию уже существующего вектора.

Output

[[5], [11], [17]]

Пример 7. Умножение матрицы на матрицу:

matrix1 = np.array([[1, 2], [3, 4], [5, 6]])
matrix2 = np.array([[1, 0], [0, 1]])
result = np.einsum("ik,kj->ij", matrix1, matrix2)
print(result)
Output

[[1, 2], [3, 4], [5, 6]]

Пример 8. Скалярное произведение векторов:

vector1 = np.array([[1, 2, 3]])
vector2 = np.array([[1, 1, 1]])
result = np.einsum("ik,jk->", vector1, vector2)
print(result)
Output

6

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

matrix1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
result = np.einsum("ii->", matrix1)
print(result)
Output

15

Пример 10. Адамарово (покомпонентное) произведение:

matrix1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
matrix2 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
result = np.einsum("ij,ij->ij", matrix1, matrix2)
print(result)

Это может показаться контринтутивно, но, как написано выше: если не понятно, что делает einsum – запиши через циклы:

result = np.zeros(matrix1.shape, dtype="int32")
for i in range(result.shape[0]):
    for j in range(result.shape[1]):
        result[i, j] += matrix1[i, j] * matrix2[i, j]
print(result)
Output

[[1, 0, 0], [0, 5, 0], [0, 0, 9]]

Пример 11. Кронекерово (внешнее) произведение векторов:

vector1 = np.array([1, 2, 3])
vector2 = np.array([1, 0, 0])
result = np.einsum("i,j->ij", vector1, vector2)
print(result)
Output

[[1, 0, 0], [2, 0, 0], [3, 0, 0]]

Пример 12. Транспонирование тензора:

A = np.array([[[0, 1], [1, 2], [2, 3]], [[1, 2], [2, 3], [3, 4]], [[2, 3], [3, 4], [4, 5]]])
result = np.einsum("ijk->jki", A)
print(result)
Output

[[[0, 1, 2], [1, 2, 3]], [[1, 2, 3], [2, 3, 4]], [[2, 3, 4], [3, 4, 5]]]

Пример 13. Произведение тензора на матрицу по третьей моде:

A = np.array([[[0, 1], [1, 2], [2, 3]], [[1, 2], [2, 3], [3, 4]], [[2, 3], [3, 4], [4, 5]]])
U = np.array([[1, 2], [2, 3]])
result = np.einsum("ijk,nk->ijn", A, U)
print(result)
Output

[[[2, 3], [5, 8], [8, 13]], [[5, 8], [8, 13], [11. 18]], [[8, 13], [11, 18], [14, 23]]]

Итоги

Конечно, einsum поставляет только дополнительный синтаксический сахар. Всегда можно использовать цепочки вложенных циклов, множество библиотечных функций (np.dot, np.outer, np.tensordot, np.transpose, np.cumsum и т.д.), и вообще не использовать einsum. Но если потратить время и понять, как она работает, то можно научиться писать гораздо более сжатый, и, не побоюсь этого слова, эффективный код.

Ссылки

Ролик с примерами einsum (ещё больше примеров).

Соглашение Эйнштейна (база)

Соглашение Эйнштейна (продвинутая часть)

Комментарии 8

    +3
    Удивительное дело, но в русскоязычном сегменте интернета почти нет материала, разъясняющего понятным языком соглашение Эйнштейна о суммировании

    Любой нормальный учебник по ОТО или ТеорМеху =)

      0
      Более того, в свое время выходила куча «Эйнштейновских сборников» на русском. Достаточно забить в поисковик — они есть. Там эти соглашения неоднократно обсуждались.
      +3
      Всегда можно использовать цепочки вложенных циклов

      Это в Python то? Возьмите достаточно большие многомерные массивы и посчитайте что-нибудь на них через "цепочки вложенных циклов". Сколько чашек чая вы успеете выпить до получения результата? Если уж пишете про функции numpy, не над описать такую глупость про "цепочки вложенных циклов", ведь кто-то может действительно подумать, что с numpy-массивами в питоне надо работать через циклы, потому что привычнее и проще, а не то что эти ваши векторизованные операции над массивами вроде всяких непонятных einsum.

        0
        Спасибо, очень интересно!
        Я с тензорами разбирался по книжке Борисенко и Тарапова, но там это подробно не рассмотрено.
          +7

          Замечу, что у функции einsum есть параметр optimize, который отвечает за то, будут ли ваши суммы считаться эффективно или как придётся. По дефолту оптимизация выключена, поэтому при неудачном стечении обстоятельств вычисления могут получиться очень неоптимальными. В примере ниже вычисляется вектор A*A*x. Без оптимизации сначала перемножаются две матрицы n x n за O(n^3), а потом уже произойдёт умножение на вектор. С оптимизацией всё будет хорошо: произойдёт два умножения матрицы на вектор и сложность будет O(n^2).


          import numpy as np
          n = 1000
          x = np.ones((n,))
          A = np.ones((n,n))
          %time _ = np.einsum("ij,jk,k->i", A,A,x)
          > CPU times: user 2.51 s, sys: 3 µs, total: 2.51 s
          > Wall time: 2.51 s
          %time _ = np.einsum("ij,jk,k->i", A,A,x, optimize=True)
          > CPU times: user 635 µs, sys: 8.01 ms, total: 8.65 ms
          > Wall time: 8.55 ms

          Можно ещё указывать optimize="optimal", тогда из всех вариантов вычисления суммы будет выбран оптимальный, в то время как optimize=True (алиас для optimize="greedy") означает использование жадного алгоритма для оптимизации.

            0

            Более того, optimize="optimal" почему-то умеет заставлять einsum работать с тензорами, у которых компоненты — произвольные объекты (например из sympy).
            Например, код


            x,y = symbols("x y")
            coords = np.array([x,y])
            g = np.identity(2)
            print(np.einsum("ij,i,j", g, coords, coords))

            выдает ошибку TypeError: invalid data type for einsum
            Но если сделать


            print(np.einsum("ij,i,j", g, coords, coords, optimize="optimal"))

            внезапно получается правильный ответ 1.0*x**2 + 1.0*y**2.

              +2

              Вообще звучит как баг.

            –1
            Жалко Эйнштейн не дожил до Python. Или хотя бы до многомерных массивов в ранних ЯП. Фиктивные индексы, индекс с запятой до стрелочки — что за дурь! Бритва Окама в помощь.

            Только полноправные пользователи могут оставлять комментарии. Войдите, пожалуйста.

            Самое читаемое