5 Best Ways to Get the Trace of a Matrix with Einstein Summation Convention in Python

by

💡 Problem Formulation: In numerical computing and linear algebra, the trace of a matrix is the sum of its diagonal elements. Utilizing the Einstein summation convention in Python can simplify operations such as finding the trace. The goal is to get the trace of a matrix, for example, given the matrix [[1, 2], [3, 4]], the desired output is 5.

Method 1: Use NumPy’s einsum function

Einstein summation (einsum) is a powerful notation that can compactly represent complex sums and products for multidimensional arrays. The np.einsum function in NumPy offers a way to compute the trace of a matrix using Einstein’s summation convention. The function specification involves passing the einsum string and the matrix. The string ‘ii’ indicates summation over the i-th diagonal elements of the matrix.

Here’s an example:

import numpy as np

matrix = np.array([[1, 2], [3, 4]])
trace = np.einsum('ii', matrix)
print(trace)

Output:

5

In this snippet, np.einsum('ii', matrix) calculates the trace by summing the elements on the diagonal (i.e., elements (0,0) and (1,1)) of the given matrix. The ‘ii’ notation specifies that operation.

Method 2: Utilizing numpy.trace() function

Although not directly involving Einstein summation, the numpy.trace() function provides a straightforward and intuitive way to compute the trace of a matrix. The function computes the sum of diagonal elements of the input array.

Here’s an example:

import numpy as np

matrix = np.array([[1, 2], [3, 4]])
trace = np.trace(matrix)
print(trace)

Output:

5

The code uses the built-in NumPy function np.trace(matrix) to directly compute the trace of the array. It’s simple and effective for quick operations where the full power of Einstein notation is not required.

Method 3: Combine numpy.diagonal() with numpy.sum()

By retrieving the diagonal elements of a matrix using numpy.diagonal() and then computing the sum with numpy.sum(), you can calculate the matrix trace. This can be seen as a more step-by-step approach in contrast to methods that calculate the trace directly.

Here’s an example:

import numpy as np

matrix = np.array([[1, 2], [3, 4]])
diagonal_elements = np.diagonal(matrix)
trace = np.sum(diagonal_elements)
print(trace)

Output:

5

The code first retrieves the diagonal elements of the matrix which is a vector, then sums over this vector to find the trace. While this method is more verbose, it explicitly demonstrates the process of finding the trace.

Method 4: Custom Einstein Summation Function

For educational purposes or to avoid NumPy’s overhead for small matrices, one might implement a custom function mimicking Einstein’s summation convention. Such a function would iterate over the indices of the matrix’s diagonal and compute the sum.

Here’s an example:

def einstein_trace(matrix):
    size = len(matrix)
    return sum(matrix[i][i] for i in range(size))

matrix = [[1, 2], [3, 4]]
trace = einstein_trace(matrix)
print(trace)

Output:

5

This user-defined function called einstein_trace(matrix) computes the trace by looping over the index of the square matrix and summing over the diagonal elements. It’s more in the spirit of the Einstein summation convention without relying on external libraries.

Bonus One-Liner Method 5: Use List Comprehension and sum() Function

For those who appreciate Python’s conciseness, we can compute the trace of a matrix in a one-liner using list comprehension combined with the built-in sum() function.

Here’s an example:

matrix = [[1, 2], [3, 4]]
trace = sum(matrix[i][i] for i in range(len(matrix)))
print(trace)

Output:

5

This one-liner extracts and sums the diagonal elements from the matrix using a straightforward list comprehension. It’s elegant and requires no additional libraries.

Summary/Discussion

  • Method 1: NumPy’s einsum. Very fast and concise. Requires understanding of Einstein notation. NumPy dependency.
  • Method 2: numpy.trace(). Extremely easy to use and read. Not directly related to Einstein summation. NumPy dependency.
  • Method 3: numpy.diagonal() with numpy.sum(). Explicit demonstration of obtaining the trace. More verbose. NumPy dependency.
  • Method 4: Custom Einstein Summation Function. No dependencies and educational. Less efficient than built-in functions.
  • Bonus Method 5: List Comprehension with sum(). Pythonic and no NumPy needed. Not as efficient with very large matrices.