How to Use NumPy’s einsum for Efficient Array Computation (5 examples)

NumPy’s einsum function is an incredibly powerful tool for executing Einstein summation convention, which can significantly optimize and speed up a wide variety of linear algebra operations. This tutorial will guide you through the ins and outs of utilizing einsum effectively within your Python code, complete with illustrative examples and practical uses.

Introduction to Einstein Summation with NumPy

Einstein summation convention is a shorthand notation that simplifies tensor algebra by conveniently expressing complex sums of products. NumPy’s einsum allows us to replicate this approach programmatically. Let’s start with the basics of einsum syntax and gradually move to more complex applications.

Basic Syntax:

import numpy as np
result = np.einsum('subscript_string', operand1, operand2, ...)

The ‘subscript_string’ is a comma-separated list defining indices of operands and the operation to perform.

Basic Operations Using einsum

Let’s begin with some simple examples to understand the primary capabilities of einsum.

Example 1: Sum of an Array

import numpy as np
array = np.array([1, 2, 3])
sum = np.einsum('i->', array)
print(sum) # Output: 6

In this example, ‘i->’ indicates sum over the i-th axis of the array.

Example 2: Element-wise Multiplication and Sum

array1 = np.array([1, 2, 3])
array2 = np.array([0, 1, 0])
dot_product = np.einsum('i,i->', array1, array2)
print(dot_product) # Output: 2

It illustrates a dot product; ‘i,i->’ signifies element-wise multiplication over the i-th indices and sum.

Advanced Operations with einsum

Moving towards more sophisticated examples, we will expand the capabilities of einsum to perform advanced linear algebra computations.

Example 3: Matrix Transposition

matrix = np.array([[1, 2], [3, 4]])
transpose = np.einsum('ij->ji', matrix)
print(transpose)
# Output:
# [[1 3]
#  [2 4]]

Here, ‘ij->ji’ represents the transposition of a 2D matrix, switching the i-th and j-th indices.

Example 4: Matrix Multiplication (Dot Product)

matrix1 = np.array([[1, 0], [0, 1]])
matrix2 = np.array([[4, 1], [2, 2]])
product = np.einsum('ij,jk->ik', matrix1, matrix2)
print(product)
# Output:
# [[4 1]
#  [2 2]]

‘ij,jk->ik’ indicates matrix multiplication, where i and k form the output indices and j sums over.

Optimizations with Einsum

Aside from simplifying operations, einsum also provides substantial performance improvements for certain types of computations. The optimization capabilities can be invoked by using the optimize argument.

Example 5: Optimizing Computational Paths

arrays = [np.random.rand(100, 100) for _ in range(10)]
multi_dot = np.einsum('ij,jk,kl,mn,np,pq,qr,rs,st,tu->iu', *arrays, optimize=True)
print(multi_dot.shape) # Output: (100, 100)

With optimize=True, NumPy automatically identifies the best order of operations, hence reducing the computational complexity.

Conclusion

This guide walked you through using NumPy’s einsum for various array operations. We learned its basic syntax, practiced common operations, and delved into advanced examples and optimizations. Employing einsum can lead to more readable, efficient, and faster code, particularly for complex tensor computations.