NumPy Einsum for Scientific Computing

Introduction

In scientific computing, it is often necessary to perform various linear algebra operations. NumPy is a popular Python library that provides efficient and convenient tools for performing such operations. One of the most powerful tools in NumPy is einsum, which stands for Einstein Summation.

In this tutorial, we will learn what einsum is and how to use it effectively in Python.

Skills Graph

Einsum Introduction

einsum is a powerful tool for performing various array operations in NumPy. It is a concise way of expressing complex array operations using a string notation. The string notation is based on the Einstein summation convention, which uses indices to denote which dimensions of arrays are being operated on. The indices are represented by letters or characters, and the operations are denoted by the characters that appear between the indices.

Open the Python Shell

Open the Python shell by typing the following command in the terminal.

python3

First Example

For example, let's consider a simple operation of multiplying two matrices A and B:

C_{ij} = \sum_k A_{ik}B_{kj}

Here, we sum over the index k, and the resulting matrix C has dimensions (i,j). Using einsum, we can express this operation as follows:

import numpy as np

## generate two random matrixs
A = np.random.rand(3, 4)
B = np.random.rand(4, 5)

## multiplying two matrices
C = np.einsum('ik,kj->ij', A, B)

The first argument to einsum is the string notation that describes the operation. In this case, we use the string 'ik,kj->ij', which tells einsum to perform the matrix multiplication and return the resulting matrix C with dimensions (i,j).

Anatomy of an Einsum String

The einsum string notation has the following general format:

output = np.einsum('input1_label1input1_label2, input2_label1input2_label2 -> output_label1output_label2', input1, input2)

• input1_label1input1_label2 and input2_label1input2_label2: These are the input tensor labels followed by their corresponding index labels. The index labels are separated by a comma. For example, ij would represent the indices i and j of the first input tensor.
• output_label1output_label2: This is the label and index specification for the output tensor. For example, ik would represent the indices i and k of the output tensor.
• input1 and input2: These are the input tensors.

The notation allows for very flexible and concise specification of tensor operations, making it a powerful tool for scientific computing.

Basic Einsum Examples

Let's look at some basic examples to get a better understanding of how einsum works.

Dot product

Suppose we have two vectors u and v, and we want to compute their dot product:

a = \sum_i u_i v_i

We can use einsum to perform this operation as follows:

## generate two random matrixs
u = np.random.rand(3)
v = np.random.rand(3)

## compute dot product
np.einsum('i,i->', u, v)

Here, the string 'i,i->' tells einsum to perform the dot product and return a scalar value.

This is equivalent to use np.dot function:

np.dot(u, v)

Transpose

Suppose we have a matrix A, and we want to compute its transpose:

B_{ij} = A_{ji}

We can use einsum to perform this operation as follows:

## generate a random matrix
A = np.random.rand(3, 4)

## transpose the matrix
np.einsum('ji->ij', A)

Here, the string 'ji->ij' tells einsum to transpose the matrix A and return the resulting matrix.

This is equivalent to use np.transpose function:

np.transpose(A)

Matrix Multiplication

Suppose we have two matrices A and B, and we want to compute their product:

C_{ij} = \sum_k A_{ik}B_{kj}

We can use einsum to perform this operation as follows:

## generate two random matrixs
A = np.random.rand(3, 4)
B = np.random.rand(4, 5)

## compute product
np.einsum('ik,kj->ij', A, B)

Here, the string 'ik,kj->ij' tells einsum to perform the matrix multiplication and return the resulting matrix with dimensions (i,j).

This is equivalent to use np.matmul function:

np.matmul(A, B)

Broadcasting

Suppose we have a matrix A and a vector u, and we want to compute the product:

B_{ij} = A_{ij}u_j

We can use einsum to perform this operation as follows:

## generate two random matrixs
A = np.random.rand(3, 4)
u = np.random.rand(4)

np.einsum('ij,j->ij', A, u)

Here, the string 'ij,j->ij' tells einsum to perform the element-wise multiplication between A and u, and return the resulting matrix.

This is equivalent to use np.multiply function:

np.multiply(A, u)

Advanced Einsum Examples

Now that we have covered the basics of einsum, let's look at some more advanced examples.

Tensor Contraction

Suppose we have a 4th-order tensor T with dimensions (i,j,k,l), and we want to contract it along the last two dimensions to obtain a 2nd-order tensor S with dimensions (i,j):

S_{ij} = \sum_{k,l} T_{ijkl}

We can use einsum to perform this operation as follows:

## generate a 4th-order tensor
T = np.random.rand(3, 4, 5, 6)

## contract
np.einsum('ijkl->ij', T)

Here, the string 'ijkl->ij' tells einsum to contract the tensor T along the last two dimensions and return the resulting matrix.

This is equivalent to use np.sum function:

np.sum(T, axis=(2, 3))

Batch Matrix Multiplications

Suppose we have a batch of n matrices A_1, ..., A_n with dimensions (m, p) and a batch of n matrices B_1, ..., B_n with dimensions (p, q), and we want to compute their batch product C_1, ..., C_n with dimensions (m, q):

C_i = A_iB_i

We can use einsum to perform this operation as follows:

n = 5
m = 3
p = 4
q = 5

## generate two matrixs
A = np.random.rand(n, m, p)
B = np.random.rand(n, p, q)

## compute batch matrix multiplication
np.einsum('nmp,npq->nmq', A, B)

Here, the string 'nmp,npq->nmq' tells einsum to perform the batch matrix multiplication and return the resulting batch of matrices.

This is equivalent to use np.matmul function:

np.matmul(A, B)

Summary

In this tutorial, we have learned what einsum is and how to use it effectively in NumPy. We have covered the basics of the einsum function, including its string notation format, and have explored several examples of how it can be used to perform various array operations such as dot products, matrix multiplications, and tensor contractions.

By mastering einsum, you can perform complex array operations in a more efficient and concise manner, saving valuable time and computational resources. So go ahead and give it a try in your next scientific computing project!