Numpy Einstein Summation

                                                            Target audience: Beginner
                                                            Estimated reading time: 3' 

Posts history
Newsletter: Geometric Learning in Python 

Many research papers use Einstein summation notation to describe mathematical concepts. Wouldn't it be great to have a similar method or operator in Python? Fortunately, NumPy provides an excellent solution for this.

Table of content

Overview

Matrix operations

Examples

       Kalman filter state equation

       Neural linear transformation

       Simplified Einstein field equation

References

 

What you will learn: How to write an elegant implementation of any tensor equation using numpy einsum

Notes: 

• Environments: Python 3.12.5,  Numpy 2.2.0
• Source code is available on GitHub [ref 1]
• To enhance the readability of the algorithm implementations, we have omitted non-essential code elements like error checking, comments, exceptions, validation of class and method arguments, scoping qualifiers, and import statement.

Overview

The NumPy einsum method [ref 2] offers two modes of application:

1. Implicit Subscripts: The output format is inferred based on the input 

    subscripts.  

   Syntax: np.einsum('subscript1, subscript2, ...', tensor1, tensor2)

   Example:   np.einsum('ji', t) computes the transpose of np.einsum('ij', t).

2. Explicit Subscripts: The output format is explicitly specified. 

    Syntax np.einsum('subscript1, subscript2, ... -> output', tensor1, tensor2)

    Example:  np.einsum('ij,jk->ik', a, b) computes the matrix multiplication of 

                    a and b.

Note: Broadcasting is supported using ellipsis notation.  

        Example:  np.einsum('...ii->...i', t) extracts the diagonal along the last two 

                         dimensions.

Matrix operations

The NumPy einsum method offers an elegant alternative to the commonly used @ operator or dedicated NumPy methods.

Matrix multiplication: einsum can be used as an alternative to the numpy operator @

a = np.array([[1.0, 0.5], [2.0, 1.5]])
b = np.array([[0.1, 2.0], [1.2, 0.5]])

output_matrix = np.einsum('ij,jk->ik', a, b)
ref_matrix = a @ b

Output:

[[0.70   2.25]

 [2.00   4.75]]

Dot product: Alternative to numpy.dot method

a = np.array([1.0, 0.5, 2.0])
b = np.array([0.1, 2.0, 1.2])

np.einsum('i,i->', a, b)
ref_dot_value = np.dot(a, b)

Output:  3.5

Matrix element-wise multiplication: Alternative to * operator

a = np.array([[1.0, 0.5], [2.0, 1.5]])
b = np.array([[0.1, 2.0], [1.2, 0.5]])

output = np.einsum('ij,ij->', a, b)
ref_output = a* b

Output: 4.25

Matrix outer product: Alternative to np.outer method

a = np.array([1.0, 0.5, 2.0, 0.7])
b = np.array([0.1, 2.0, 1.2])

output = np.einsum('i,j->ij', a, b)
ref_output = np.outer(a, b)

Output:

[[0.10   2.00   1.20 ]

 [0.05  1.00    0.60 ]

 [0.20  4.00    2.40 ]

 [0.07  1.40    0.84]]

Matrix Transpose: Alternative to .T operator

a = np.array([[1.0, 0.5], [2.0, 1.5]])

output = np.einsum('ij->ji', a)
ref_output = a.T

Output:

[[1.  2. ]

 [0.5 1.5]]

Matrix trace: Alternative to np.trace method

a = np.array([[1.0, 0.5], [2.0, 1.5]])

output = np.einsum('ii->', a)
ref_output = np.trace(a)

Output: 2.5

Examples

Kalman filter state equation

The state transition equation of the Kalman filter [ref 3, 4] is defined as

\[ \widetilde{x}_{k/k-1}=A_{k}.\widetilde{x}_{k-1/k-1} + B_{k}.u_{k} + w \]

A = np.array([[1.0, 0.1], [0.0, 1.0]])   # State transition matrix
B = np.array([[0.5], [1.0]])                 # Control input matrix
x_k_1 = np.array([2.0, 1.0])              # State vector at time k-1
u = np.array([0.1])                             # Control input
w = np.array([0.05, 0.02])   
 
# Using numpy matrix operator
x_k = A @ x_k_1 + B @ u + w

# Using Einstein summation method
x_k = np.einsum('ij,j->i', A, x_k_1) + np.einsum('ij,j->i', B, u) + w

Output:

[2.20, 1.12]

Neural linear transformation

The linear transformation of a layer of multi-layer perceptron is defined as \[ \textbf{y} = W.\textbf{x} + b \] W is the weight matrix,

x is the input vector,

b is the bias vector,

y is the output vector.

W = np.array([[0.20,  0.80, -0.50], 
                       [0.50, -0.91,  0.26], 
                       [-0.26, 0.27,  0.17]])  # Shape (3, 3)
x = np.array([1, 2, 3])  # Input vector, shape (3,)
b = np.array([2.0, 3.0, 0.5])  # Bias vector, shape (3,)

# Using @ operator
y = W @ x + b

# Using Einstein summation
y_einsum = np.einsum('ij,j->i', W, x) + b

Output:

[2.30  2.46 1.29]

Simplified Einstein field equation

The Einstein field equations (EFE) in general relativity are complex tensor equations that describe the relationship between the geometry of spacetime and the matter and energy within it: \[ G_{\mu \nu }= R_{\mu \nu} - \frac{1}{2}Rg_{\mu \nu} = \frac{8\pi G}{c^{4}} T_{\mu \nu} \]

# Define the Riemannian metric tensor (4-dimension)
g = np.array([[-1, 0, 0, 0],
                     [  0, 1, 0, 0],  
                     [  0, 0, 1, 0],
                     [  0, 0, 0, 1]])  # Shape (4, 4)

# Ricci curvature tensor
R_mu_nu = np.random.rand(4, 4)  # Shape (4, 4)

# R: trace of Ricci tensor as scalar curvature
R = np.einsum('ii', np.linalg.inv(g) @ R_mu_nu)  #

# Compute Einstein tensor G_mu_nu
G_mu_nu = R_mu_nu - 0.5 * R * g

References

[1] geometriclearning/geometry/ein_sum

[2] Numpy einsum documentation

[3] Distributed Recursive Kalman Filter on Large Datasets

[4] Limitation of Linear Kalman Filter

-------------
Patrick Nicolas has over 25 years of experience in software and data engineering, architecture design and end-to-end deployment and support with extensive knowledge in machine learning. 
He has been director of data engineering at Aideo Technologies since 2017 and he is the author of "Scala for Machine Learning", Packt Publishing ISBN 978-1-78712-238-3 and Geometric Learning in Python Newsletter on LinkedIn.