Array axis summations with Einstein summation convention in Python

The einsum() method evaluates the Einstein summation convention on operands. Using the Einstein summation convention, many common multi-dimensional, linear algebraic array operations can be represented in a simple fashion. This method provides flexibility to compute array operations including axis summations by specifying subscript labels.

For array axis summations with Einstein summation convention, use the numpy.einsum() method. The first parameter is the subscript string that specifies the summation pattern, and the second parameter is the input array.

Syntax

numpy.einsum(subscripts, *operands)

Parameters

• subscripts − String specifying the subscripts for summation as comma separated list of subscript labels
• operands − Input arrays for the operation

Basic Array Axis Summation

Let's create a 2D array and perform axis summations using different subscript patterns ?

import numpy as np

# Create a 2D array
arr = np.arange(12).reshape(3, 4)
print("Original Array:")
print(arr)
print("Shape:", arr.shape)

Original Array:
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
Shape: (3, 4)

Sum Along Different Axes

The subscript pattern determines which axis to sum over ?

import numpy as np

arr = np.arange(12).reshape(3, 4)

# Sum along axis 1 (rows): 'ij->i'
print("Sum along axis 1 (sum each row):")
print(np.einsum('ij->i', arr))

# Sum along axis 0 (columns): 'ij->j' 
print("\nSum along axis 0 (sum each column):")
print(np.einsum('ij->j', arr))

# Sum all elements: 'ij->'
print("\nSum all elements:")
print(np.einsum('ij->', arr))

Sum along axis 1 (sum each row):
[ 6 22 38]

Sum along axis 0 (sum each column):
[12 15 18 21]

Sum all elements:
66

Understanding the Subscript Pattern

The subscript pattern follows the format 'input_indices->output_indices' ?

import numpy as np

arr = np.arange(16).reshape(4, 4)
print("4x4 Array:")
print(arr)

# 'ij->i' means: for each row i, sum over all columns j
print("\n'ij->i' - Sum each row:")
print(np.einsum('ij->i', arr))

# 'ij->j' means: for each column j, sum over all rows i  
print("\n'ij->j' - Sum each column:")
print(np.einsum('ij->j', arr))

4x4 Array:
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]

'ij->i' - Sum each row:
[ 6 22 38 54]

'ij->j' - Sum each column:
[24 28 32 36]

Comparison with Standard NumPy Methods

3D Array Example

Einstein summation works with higher dimensional arrays too ?

import numpy as np

# Create a 3D array (2x3x4)
arr_3d = np.arange(24).reshape(2, 3, 4)
print("3D Array shape:", arr_3d.shape)

# Sum along the last axis: 'ijk->ij'
print("\nSum along last axis (k):")
print(np.einsum('ijk->ij', arr_3d))

# Sum along middle axis: 'ijk->ik'
print("\nSum along middle axis (j):")
print(np.einsum('ijk->ik', arr_3d))

3D Array shape: (2, 3, 4)

Sum along last axis (k):
[[ 6 22 38]
 [54 70 86]]

Sum along middle axis (j):
[[36 42 48 54]
 [72 78 84 90]]

Conclusion

The numpy.einsum() method provides a flexible way to perform axis summations using Einstein notation. The subscript pattern 'ij->i' sums along the second dimension, while 'ij->j' sums along the first dimension, offering an alternative to standard NumPy sum operations.