Get the trace of a matrix with Einstein summation convention in Python

The einsum() method evaluates the Einstein summation convention on operands. Using the Einstein summation convention, many common multi-dimensional, linear algebraic array operations can be represented in a simple fashion.

To get the trace of a matrix with Einstein summation convention, use the numpy.einsum() method. The trace is the sum of diagonal elements, which can be computed using the subscript 'ii'.

Syntax

numpy.einsum(subscripts, *operands)

Parameters:

• subscripts − String specifying the subscripts for summation
• operands − Input arrays for the operation

Basic Example

Let's create a 4x4 matrix and calculate its trace using einsum() ?

import numpy as np

# Create a 4x4 matrix
arr = np.arange(16).reshape(4, 4)
print("Matrix:")
print(arr)

# Calculate trace using einsum
trace_result = np.einsum('ii', arr)
print(f"\nTrace using einsum: {trace_result}")

# Verify with numpy's trace function
trace_verify = np.trace(arr)
print(f"Trace using np.trace: {trace_verify}")

Matrix:
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]

Trace using einsum: 30
Trace using np.trace: 30

How Einstein Summation Works for Trace

The subscript 'ii' tells NumPy to sum over elements where both indices are equal (diagonal elements). This is equivalent to arr[0,0] + arr[1,1] + arr[2,2] + arr[3,3].

import numpy as np

# Create a 3x3 matrix for clearer visualization
matrix = np.array([[1, 2, 3],
                   [4, 5, 6], 
                   [7, 8, 9]])

print("3x3 Matrix:")
print(matrix)

# Manual calculation of diagonal sum
manual_trace = matrix[0,0] + matrix[1,1] + matrix[2,2]
print(f"\nManual trace calculation: {matrix[0,0]} + {matrix[1,1]} + {matrix[2,2]} = {manual_trace}")

# Using einsum
einsum_trace = np.einsum('ii', matrix)
print(f"Einsum trace: {einsum_trace}")

3x3 Matrix:
[[1 2 3]
 [4 5 6]
 [7 8 9]]

Manual trace calculation: 1 + 5 + 9 = 15
Einsum trace: 15

Comparison with Other Methods

Advanced Example

Einstein summation can also calculate traces of multiple matrices simultaneously ?

import numpy as np

# Create a 3D array containing multiple 3x3 matrices
matrices = np.random.randint(1, 10, size=(2, 3, 3))
print("Multiple matrices:")
print(matrices)

# Calculate trace for each matrix using einsum
traces = np.einsum('aii->a', matrices)
print(f"\nTraces: {traces}")

# Verify with individual calculations
for i, matrix in enumerate(matrices):
    individual_trace = np.trace(matrix)
    print(f"Matrix {i} trace: {individual_trace}")

Multiple matrices:
[[[8 2 7]
  [6 4 1]
  [5 9 3]]

 [[2 8 6]
  [7 1 9]
  [4 5 3]]]

Traces: [15  6]
Matrix 0 trace: 15
Matrix 1 trace: 6

Conclusion

The numpy.einsum() method with subscript 'ii' provides a flexible way to calculate matrix trace using Einstein summation convention. While np.trace() is more direct for simple cases, einsum() offers greater flexibility for complex operations involving multiple matrices.