Using np.einsum() to avoid explicit loops in NumPy

Introduction

One of the lesser-known, yet incredibly powerful tools in NumPy that can help improve the efficiency of your numerical computations is np.einsum(). This function provides a succinct way to perform operations on arrays of any dimension and is based on Einstein summation convention, from which it derives its name. In this tutorial, we’re going to delve deep into the workings of np.einsum(), demonstrating its use through a series of examples from basic to advanced.

Before diving into the practical applications, let’s understand what np.einsum() brings to the table. Traditional operations with NumPy arrays often involve explicit loops or broadcasting which can be slower or less memory efficient. np.einsum() allows you to specify complex operations in a more optimal way, thereby simplifying multidimensional arithmetic and potentially leading to significant speedups.

Basic Usage of np.einsum()

Note: Before running the samples, make sure you have NumPy installed and imported with import numpy as np.

At its core, np.einsum() takes a string that describes the subscripts for summation and the operands on which to perform the operation. The subscripts are followed by arrow (->) and the desired output’s subscript, if required.

Example 1: Summing Elements of an Array

import numpy as np

a = np.array([1, 2, 3])
sum_a = np.einsum('i->', a)
print(sum_a)

Output:

6

This example simply sums up all the elements of array a by specifying a subscript 'i' and the result being a scalar hence no return subscript is specified ('->' alone).

Vector and Matrix Multiplication

Example 2: Dot Product of Two Vectors

import numpy as np

v1 = np.array([1, 2, 3])
v2 = np.array([4, 5, 6])
dot_product = np.einsum('i,i->', v1, v2)
print(dot_product)

Output:

32

Here, we used the same subscript 'i' for both vectors indicating element-wise multiplication followed by a summation across that axis.

Example 3: Matrix Multiplication

import numpy as np

m1 = np.array([[1, 2], [3, 4]])
m2 = np.array([[5, 6], [7, 8]])
product = np.einsum('ij,jk->ik', m1, m2)
print(product)

Output:

[[19 22]
 [43 50]]

This represents a standard matrix multiplication. The result 'ik' indicates that we’re taking rows from m1 and columns from m2.

Advanced Operations with np.einsum()

As we become comfortable with the subscript notation, we can perform even more complex array manipulations.

Example 4: Tensor Contraction

import numpy as np

a = np.arange(27).reshape(3, 3, 3)
b = np.arange(3)
contraction = np.einsum('ijk,k->ij', a, b)
print(contraction)

Output:

[[ 15  42  69]
 [ 96 123 150]
 [177 204 231]]

This operation is similar to matrix multiplication, but it involves a 3-dimensional tensor a. The vector b is contracted over the last axis of a, resulting in a 2-dimensional array.

Example 5: Transposition and Diagonal

import numpy as np

m1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
transpose_diag = np.einsum('ii->i', m1)
print(transpose_diag)

Output:

[1 5 9]

Instead of using np.diagonal(), einsum provides a diagonal operation along with a transpose if needed. The repeated subscript 'ii' indicates that we’re selecting the diagonal of our matrix.

The true power of np.einsum() is realized in converting seemingly complex operations into concise and efficient executions. Advanced algorithms can be implemented with einsum, benefiting from its capacity to handle broadcasting, permuting, and other array operations cohesively.

Optimization Tips

Using np.einsum() efficiently involves knowing a few optimization strategies.

• Be mindful of the order of operations as it can affect performance; use intermediate results if beneficial.
• Starting with NumPy version 1.12.0, np.einsum() includes an optimize flag that attempts to identify the optimal order of operations on its own.
• For large scale problems, consider using libraries that incorporate automatic optimization for tensor operations, such as Dask or TensorFlow.

Conclusion

In the landscape of high-performance computing with Python, np.einsum() offers a flexible and powerful way to express complex array operations. Through the examples we’ve discussed, you can now harness this tool to avoid explicit loops and make your computations faster and more readable. Dive into the documentation for more details and embrace the efficiency it offers to both your code and your coding mindset.