What are the meanings of subscripts passed to numpy.einsum()?

Question

I am trying to understand a python code, which uses numpy.einsum() to convert a 4-dimensional numpy array, A, to 2- or 3-dimensional arrays. The subscripts passed to numpy.einsum() are as follows:

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)

etc. Following the answers of (Understanding NumPy's einsum) and (Python - Sum 4D Array) I tried to use numpy.sum() to understand the meanings of the above subscripts, for example, Mat1 = np.sum(A, axis=(0,3)) but I could not reproduce the results, which I get with numpy.einsum(). Can someone please explain how these subscripts are interpreted in numpy.einsum() ?

Answer 1

I advise you to read Einstein notation on Wikipedia.

Here's a short answer to your question:

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

means:

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

Short explanation:
aabb means the indexes and their equality (see A[a, a, b, b]);
->ab means the shape is (max_a, max_b) and you don't need two have sum on these two indexes. (if their were c also then you should sum everything by c as it is not presented after ->)

---

Other your examples:

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

Some code to check that it is actually true:

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