I delve into a diverse range of topics, spanning programming languages, machine learning, data engineering tools, and DevOps. Our articles are enriched with practical code examples, ensuring their applicability in real-world scenarios.
Friday, January 3, 2025
Thursday, January 2, 2025
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
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]])
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 matrixB = np.array([[0.5], [1.0]]) # Control input matrixx_k_1 = np.array([2.0, 1.0]) # State vector at time k-1u = np.array([0.1]) # Control inputw = np.array([0.05, 0.02])# Using numpy matrix operatorx_k = A @ x_k_1 + B @ u + w# Using Einstein summation methodx_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 @ operatory = W @ x + b# Using Einstein summationy_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
-------------
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.
Subscribe to:
Posts (Atom)