传递给 numpy.einsum() 的下标是什么意思？

Question

我试图理解一个 python 代码，它使用 numpy.einsum() 将 4 维 numpy 数组 A 转换为 2 维或 3 维数组。传递给numpy.einsum()的下标如下：

Mat1 = np.einsum('aabb->ab', A) 

Mat2 = np.einsum('abab->ab', A)

Mat3 = np.einsum('abba->ab', A) 

T1 = np.einsum('abcb->abc' A)

T2 = np.einsum('abbc->abc', A)

等等。根据 (Understanding NumPy's einsum) 和 (Python - Sum 4D Array) 的答案，我尝试使用 numpy.sum() 来理解上述下标的含义，例如 Mat1 = np.sum(A, axis=(0,3)) 但我无法重现结果，我得到了 @ 987654329@。 有人可以解释一下numpy.einsum() 中这些下标是如何解释的吗？

Answer 1

我建议你阅读Einstein notation on Wikipedia。

以下是对您问题的简短回答：

np.einsum('aabb->ab', A)

意思是：

res = np.empty((max_a, max_b), dtype=A.dtype)
for a in range(max_a):
  for b in range(max_b):
    res[a, b] = A[a, a, b, b]
return res

简短说明：
aabb 表示索引及其相等性（参见A[a, a, b, b]）；
->ab 表示形状为(max_a, max_b)，您不需要两个对这两个求和索引。 （如果他们也是c，那么你应该用c对所有内容求和，因为->之后不会出现）

---

你的其他例子：

np.einsum('abab->ab', A)

# Same as (by logic, not by actual code)

res = np.empty((max_a, max_b), dtype=A.dtype)
for a in range(max_a):
  for b in range(max_b):
    res[a, b] = A[a, b, a, b]
return res

np.einsum('abba->ab', A) 

# Same as (by logic, not by actual code)

res = np.empty((max_a, max_b), dtype=A.dtype)
for a in range(max_a):
  for b in range(max_b):
    res[a, b] = A[a, b, b, a]
return res

np.einsum('abcb->abc', A)

# Same as (by logic, not by actual code)

res = np.empty((max_a, max_b, max_c), dtype=A.dtype)
for a in range(max_a):
  for b in range(max_b):
    for c in range(max_c):
      res[a, b, c] = A[a, b, c, b]
return res

np.einsum('abbc->abc', A)

# Same as (by logic, not by actual code)

res = np.empty((max_a, max_b, max_c), dtype=A.dtype)
for a in range(max_a):
  for b in range(max_b):
    for c in range(max_c):
      res[a, b, c] = A[a, b, b, c]
return res

一些代码来检查它是否真的是真的：

import numpy as np


max_a = 2
max_b = 3
max_c = 5

shape_1 = (max_a, max_b, max_c, max_b)
A = np.arange(1, np.prod(shape_1) + 1).reshape(shape_1)

print(A)
print()
print(np.einsum('abcb->abc', A))
print()

res = np.empty((max_a, max_b, max_c), dtype=A.dtype)
for a in range(max_a):
  for b in range(max_b):
    for c in range(max_c):
      res[a, b, c] = A[a, b, c, b]

print(res)
print()