Python에서 아인슈타인 합 규약(einsum)으로 행렬의 대각합(Trace) 구하기

NumPy의 einsum() 메서드는 피연산자(operand)에 대해 아인슈타인 합 규약(Einstein summation convention)을 평가합니다. 이 규약을 활용하면 다차원 배열에 대한 다양한 선형대수 연산을 매우 간결한 방식으로 표현할 수 있습니다.

암시적 모드(implicit mode)에서 einsum은 이러한 값을 자동으로 계산합니다. 반면 명시적 모드(explicit mode)에서는 특정 첨자(subscript) 레이블에 대한 합산을 비활성화하거나 강제함으로써, 고전적인 아인슈타인 합 연산으로 분류되지 않는 배열 연산까지 유연하게 수행할 수 있습니다.

아인슈타인 합 규약으로 행렬의 대각합(trace)을 구하려면 Python에서 numpy.einsum() 메서드를 사용하면 됩니다. 첫 번째 매개변수는 첨자(subscript)로, 합산에 사용할 첨자 레이블을 쉼표로 구분한 목록으로 지정합니다. 두 번째 매개변수는 피연산자(operands)로, 연산 대상이 되는 배열입니다. 행렬의 대각합을 구할 때는 'ii'처럼 동일한 첨자를 지정하여, 행 인덱스와 열 인덱스가 같은 대각선 요소들만 모두 더하게 됩니다.

단계별 진행

먼저 필요한 라이브러리를 임포트합니다.

import numpy as np

arange()와 reshape() 메서드를 사용해 NumPy 배열을 생성합니다.

arr = np.arange(16).reshape(4, 4)

생성된 배열을 출력해 확인합니다.

print("Our Array...\n", arr)

배열의 차원을 확인합니다.

print("\nDimensions of our Array...\n", arr.ndim)

배열의 데이터 타입을 확인합니다.

print("\nDatatype of our Array object...\n", arr.dtype)

배열의 형태(shape)를 확인합니다.

print("\nShape of our Array object...\n", arr.shape)

아인슈타인 합 규약으로 행렬의 대각합을 구하기 위해 numpy.einsum() 메서드를 사용합니다.

print("\nResult (trace)...\n", np.einsum('ii', arr))

전체 예제 코드

import numpy as np

# arange()와 reshape() 메서드로 NumPy 배열 생성
arr = np.arange(16).reshape(4, 4)

# 배열 출력
print("Our Array...\n", arr)

# 차원 확인
print("\nDimensions of our Array...\n", arr.ndim)

# 데이터 타입 확인
print("\nDatatype of our Array object...\n", arr.dtype)

# 형태(shape) 확인
print("\nShape of our Array object...\n", arr.shape)

# 아인슈타인 합 규약으로 행렬의 대각합 계산
print("\nResult (trace)...\n", np.einsum('ii', arr))

실행 결과

Our Array...
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]

Dimensions of our Array...
2

Datatype of our Array object...
int64

Shape of our Array object...
(4, 4)

Result (trace)...
30

위 예제에서 대각합은 주대각선 요소인 0 + 5 + 10 + 15 = 30으로 계산됩니다. 즉, np.einsum('ii', arr)는 행과 열 인덱스가 같은 요소들을 모두 더한 값, 바로 행렬의 대각합(trace)을 반환합니다.

참고로 이 결과는 np.trace(arr)와 동일하지만, einsum을 사용하면 대각합뿐 아니라 전치, 내적, 외적 등 다양한 텐서 연산을 하나의 통일된 문법으로 처리할 수 있다는 장점이 있습니다.