np.einsum() for efficient computation in NumPy

Introduction

np.einsum() helps you do math with arrays in a fast and clear way. It lets you write complex calculations simply and runs them efficiently.

When you want to multiply and sum elements of arrays in custom ways.

When you need to do matrix multiplication or dot products with less code.

When you want to avoid slow loops and speed up calculations.

When you want to combine multiple array operations into one step.

When you want to write clear code for complex math with arrays.

Syntax

NumPy

np.einsum(subscripts, *operands, optimize=False)

subscripts is a string that tells how to multiply and sum array elements.

operands are the arrays you want to calculate with.

Examples

Calculates the dot product of two 1D arrays a and b.

NumPy

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

Performs matrix multiplication of 2D arrays A and B.

NumPy

np.einsum('ij,jk->ik', A, B)

Extracts the diagonal elements of a square matrix M.

NumPy

np.einsum('ii->i', M)

Sums over the second axis of a 3D array T.

NumPy

np.einsum('ijk->ik', T)

Sample Program

This program shows how to use np.einsum() to do a dot product and matrix multiplication. It prints the results clearly.

NumPy

import numpy as np

# Define two 1D arrays
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

# Dot product using np.einsum
result_dot = np.einsum('i,i->', a, b)

# Define two 2D arrays for matrix multiplication
A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])

# Matrix multiplication using np.einsum
result_matmul = np.einsum('ij,jk->ik', A, B)

# Print results
print('Dot product:', result_dot)
print('Matrix multiplication:\n', result_matmul)

Important Notes

Use simple letters like i, j, k to label array axes in the subscripts.

np.einsum() can replace many common numpy functions like dot, sum, and transpose.

It often runs faster because it avoids creating temporary arrays.

Summary

np.einsum() lets you write complex array math clearly and efficiently.

It uses a string to describe how to multiply and sum array elements.

Great for dot products, matrix multiplication, and custom sums without slow loops.