Extract the diagonal of a matrix with Einstein summation convention in Python

The Einstein summation convention provides a concise way to express array operations. The numpy.einsum() method implements this convention, allowing us to extract matrix diagonals using subscript notation like 'ii->i'.

Understanding Einstein Summation for Diagonals

The subscript 'ii->i' means we take elements where both indices are equal (diagonal elements) and output them as a 1D array. The repeated index i on the left indicates we want elements at positions (0,0), (1,1), (2,2), etc.

Basic Example

Here's how to extract diagonal elements from a 4x4 matrix ?

import numpy as np

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

# Extract diagonal using Einstein summation
diagonal = np.einsum('ii->i', arr)
print("\nDiagonal elements:")
print(diagonal)

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

Diagonal elements:
[ 0  5 10 15]

How It Works

The notation 'ii->i' can be broken down as follows:

• ii on the left: Input array with indices i,i (same index repeated)
• ->: Arrow separating input from output specification
• i on the right: Output array with single index i

Comparison with Traditional Methods

Let's compare einsum with NumPy's built-in diag() function ?

import numpy as np

matrix = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# Method 1: Using einsum
diagonal_einsum = np.einsum('ii->i', matrix)

# Method 2: Using diag()
diagonal_diag = np.diag(matrix)

print("Matrix:")
print(matrix)
print("\nUsing einsum:", diagonal_einsum)
print("Using diag():", diagonal_diag)
print("Results match:", np.array_equal(diagonal_einsum, diagonal_diag))

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

Using einsum: [1 5 9]
Using diag(): [1 5 9]
Results match: True

Advanced Example with Different Matrix Types

Einstein summation works with different data types and matrix sizes ?

import numpy as np

# Float matrix
float_matrix = np.random.rand(3, 3).round(2)
print("Float Matrix:")
print(float_matrix)
print("Diagonal:", np.einsum('ii->i', float_matrix))

# Complex matrix
complex_matrix = np.array([[1+2j, 3+4j], [5+6j, 7+8j]])
print("\nComplex Matrix:")
print(complex_matrix)
print("Diagonal:", np.einsum('ii->i', complex_matrix))

Float Matrix:
[[0.37 0.95 0.73]
 [0.6  0.16 0.06]
 [0.7  0.02 0.97]]
Diagonal: [0.37 0.16 0.97]

Complex Matrix:
[[1.+2.j 3.+4.j]
 [5.+6.j 7.+8.j]]
Diagonal: [1.+2.j 7.+8.j]

Key Parameters

Conclusion

The numpy.einsum() method with 'ii->i' notation provides an elegant way to extract matrix diagonals. While equivalent to np.diag(), einsum offers more flexibility for complex array operations and follows mathematical notation conventions.