Advanced Array Manipulation Lesson

NumPy Broadcasting

9 min to complete · By Gilad Gressel

Broadcasting in NumPy is a set of rules by which ufuncs operate on arrays of different sizes and/or dimensions. It's sort of the magic that makes everything work under the hood, and it can be a bit complicated. I'd like you to walk away from this lesson having a reasonable idea how it works but not at all having any mastery. Note that there is a trend here - you should get familiar but only gain expertise when you know you need it.

NumPy Broadcasting with a Single Number and Array

The simplest example of broadcasting is when you have a single number and an array. In this case, the single number is broadcast to all the elements of the array. Let's see this in action.

import numpy as np
# NumPy Broadcasting is a powerful mechanism that  
# allows NumPy to work with arrays of different 
# shapes when performing arithmetic operations.
# this is the most basic example:

a = np.array([1, 2, 3])
print(a)
print()

# This will add 5 to each element of a
result = a + 5  
print(result)
[1 2 3] 
[6 7 8]

In the above example a naive implementation would be to do a for loop and add each element by 5. However, what NumPy does under the hood is actually create an array of 5s that is the same size as the original array and then adds the two arrays together. This is much faster than a python for loop (due to the underlying C routines) and is one of the reasons why NumPy is so fast.

Here is a visual representation of what broadcasting is doing.

a = np.array([1.0, 2.0, 3.0])
b = 2.0
a * b
>>>array([2.,  4.,  6.])
NumPy Broadcasting

Image Source

Here what is happening is that we are actually creating a new array of 2.0s that is the same size as the original array and then multiplying the two arrays together with elementwise multiplication.

Colorful illustration of a light bulb

Note: NumPy does not actually create copies of the scalar value, it just uses the same value over and over again. In a sense, it is broadcasting the scalar value to the array.

This is the fundamental idea of NumPy broadcasting. It gets a lot more complicated because of the rules.

NumPy Broadcasting Rules

Let's go through the rules located on the NumPy website.

The general rule is that we compare the two arrays and start from the right side and work our way left. If the dimensions are the same or one of the dimensions is 1 then we are good. If the dimensions are different and neither is 1 then we can't broadcast and we get an error.

Here are some examples of broadcasting working:

a = np.array([[ 0.0,  0.0,  0.0],
              [10.0, 10.0, 10.0],
              [20.0, 20.0, 20.0],
              [30.0, 30.0, 30.0]])
print(a.shape)
print()

b = np.array([1.0, 2.0, 3.0])
print(b.shape)
# This will add b to each row of a, it works 
# because b has the same number of columns as a
print(a + b)  
(4, 3)
    
(3,)
[[ 1.  2.  3.]
 [11. 12. 13.]
 [21. 22. 23.]
 [31. 32. 33.]]
NumPy Broadcasting

Image Source

# What happens if b has more elements than a?
b = np.array([1.0, 2.0, 3.0, 4.0]) 
# now be has 4 elements
print(a+b)
---------------------------------------------------------------------------

ValueError                                Traceback (most recent call last)

/home/amar/codingnomads/numpy_pandas_module/10_broadcasting.ipynb Cell 8 line 3
    <a href='vscode-notebook-cell://wsl%2Bubuntu-20.04/home/amar/codingnomads/numpy_pandas_module/10_broadcasting.ipynb#X23sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a> # What happens if b has more elements than a?
    <a href='vscode-notebook-cell://wsl%2Bubuntu-20.04/home/amar/codingnomads/numpy_pandas_module/10_broadcasting.ipynb#X23sdnNjb2RlLXJlbW90ZQ%3D%3D?line=1'>2</a> b = np.array([1.0, 2.0, 3.0, 4.0]) # now be has 4 elements
    ----> <a href='vscode-notebook-cell://wsl%2Bubuntu-20.04/home/amar/codingnomads/numpy_pandas_module/10_broadcasting.ipynb#X23sdnNjb2RlLXJlbW90ZQ%3D%3D?line=2'>3</a> print(a+b)


ValueError: operands could not be broadcast together with shapes (4,3) (4,) 

NumPy Broadcasting

Image Source

Here we see the mismatch makes it impossible to broadcast these two arrays together.

Broadcast Two Arrays When Both Must be Stretched

Now, note it is possible to broadcast two arrays together when both need to be "stretched" to match. This is a bit more complicated but here is an example of that.

# make a into a column vector
a = np.array([0.0, 10.0, 20.0, 30.0]).reshape(-1,1) 
print(a) 
# a is a column vector now
print(a.shape) 
print()

b = np.array([1.0, 2.0, 3.0])
print(b)   
# b is a row vector
print(b.shape) 
print()

print(a + b)
[[ 0.]
 [10.]
 [20.]
 [30.]]
(4, 1)
    
[1. 2. 3.]
(3,)
    
[[ 1.  2.  3.]
 [11. 12. 13.]
 [21. 22. 23.]
 [31. 32. 33.]]

The above broadcasting works by stretching both a and b

NumPy Broadcasting

Image Source

Here we see that a is stretched horizontally to match b, and b is stretching vertically to match a and then the two arrays are broadcast together.

That is all we will cover for NumPy broadcasting now. If you hit a broadcast error chances are you will need to add a newaxis somewhere to make things match up, or conversely squeeze extra dimensions out. In general, adding dimensions is more common before or during broadcasting to align the shapes of arrays, while squeezing dimensions often comes into play after certain operations have been performed to clean up the result for further processing.

Summary: NumPy Broadcasting

  • NumPy Broadcasting allows operations between arrays of different shapes by extending smaller arrays to match larger ones.
  • Broadcasting rules: Start comparing dimensions from the right (last dimension). If dimensions are equal, or one is 1, broadcasting can occur. If they're different and neither is 1, you encounter an error.
  • Scalars are also broadcasted to match the size of arrays during arithmetic operations.
  • For arrays with mismatched shapes, reshaping or adding new dimensions can often resolve broadcasting issues.