NumPy Broadcasting | by Pan Cretan | Medium
Definitions, rules and examples
NumPy offers fast calculations via vectorisation that avoids the use of slow Python loops. Vectorisation is also available when using binary ufuncs, such as addition or multiplication, with the added benefit that arrays do not need to have the same shape. Operations with arrays of different shape is known as broadcasting and can be particularly confusing, especially with multidimensional arrays, or when both arrays need to be stretched.
There are many examples and tutorials available, but I find most useful to approach the matter by thinking, and actually memorising, the broadcasting rules. It is then easier to think about any given use case and write the code without relying on trial-and-error.
There are two data analysis and data science books that I would highly recommend. Both of them contain a small section on broadcasting.
Python for Data Analysis by Wes McKinney contains he following broadcasting rule:
Two arrays are compatible for broadcasting if for each trailing dimension (i.e., starting from the end) the axis lengths match or if either of the lengths is 1. Broadcasting is then performed over the missing or length 1 dimensions.
Python Data Science Handbook by Jake VanderPlas contains a more verbose set of broadcasting rules:
Rule 1: If the two arrays differ in their number of dimensions, the shape of the one with fewer dimensions is padded with ones on its leading (left) side.
Rule 2: If the shape of the two arrays does not match in any dimension, the array with shape equal to 1 in that dimension is stretched to match the other shape.
Rule 3: If in any dimension the sizes disagree and neither is equal to 1, an error is raised.
I find the second set of rules easier to follow. The examples below will be using these rules.
Perhaps one of the simplest examples of broadcasting, and a typical pattern, is to subtract the column mean from each column. Following this operation the column means will become numerically equally to zero
a = np.arange(12).reshape(4, 3)
means_columns = a.mean(axis=0)
res = a - means_columns
print('original array', a, sep='\n')
print('.. column means', a.mean(axis=0), sep='\n')
print('demeaned array', res, sep='\n')
print('.. column means', res.mean(axis=0), sep='\n')
that prints
original array
[[ 0 1 2]
[ 3 4 5]
[ 6 7 8]
[ 9 10 11]]
.. column means
[4.5 5.5 6.5]
demeaned array
[[-4.5 -4.5 -4.5]
[-1.5 -1.5 -1.5]
[ 1.5 1.5 1.5]
[ 4.5 4.5 4.5]]
.. column means
[0. 0. 0.]
Let’s see what is going on here. Computing the means of the columns with a.mean(axis=0) produces an one-dimensional array with shape (3,). The two arrays involved in the subtraction differ in shape and hence means_columnsthat has fewer dimensions, will have its shape padded with 1 on the left according to rule 1. Hence, behind the scenes, means_columns will be reshaped to (1, 3). Then according to rule 2 means_columns will be stretched along axis 0 so that its shape becomes (3, 3) to match the shape of a.
In addition to using the rules to predict how the array with fewer dimensions will be stretched we can also use np.broadcast_to that returns a read-only view on the original array with the given shape. The view may be non-contiguous and it is possible that different elements refer to the same memory address
means_columns_bc = np.broadcast_to(means_columns, a.shape)
print(means_columns_bc)
print('base', means_columns_bc.base, sep='\n')
print('strides', means_columns_bc.strides, sep='\n')
that prints
[[4.5 5.5 6.5]
[4.5 5.5 6.5]
[4.5 5.5 6.5]
[4.5 5.5 6.5]]
base
[4.5 5.5 6.5]
strides
(0, 8)
We can see that the base is the original means array (hence it is a view), whilst the stride along the first axis is 0, which means that different elements of the same column refer to the same memory location (see here for a primer on NumPy internals). NumPy is truly optimising memory usage wherever it can!
What if we wanted to demean the rows? The mean of the rows can be computed readily with a.mean(axis=1) that will return a (4,) array. Padding its shape with 1 on the left, means that the array will become (1, 4). According to rule 3 the last dimension of the two arrays disagree and neither is 1. This means that broadcasting will not take place. We could also anticipate this becausenp.broadcast_to(a.mean(axis=1), a.shape) raises an exception informing us that broadcast could not produce the requested shape (4, 3). The incompatibility of shapes can also be seen by executing np.broadcast_shapes(a.shape, a.mean(axis=1).shape) that also raises an exception explaining the shape mismatch. Demeaning the rows can be done by reshaping the row means into a (4, 1) array using either a.mean(axis=1).reshape(-1, 1) or a.mean(axis=1)[:, np.newaxis]
means_rows = a.mean(axis=1)
res = a - means_rows.reshape(-1, 1) # or res = a - means_rows[:, np.newaxis]
print('original array', a, sep='\n')
print('.. row means', a.mean(axis=1), sep='\n')
print('demeaned array', res, sep='\n')
print('.. row means', res.mean(axis=1), sep='\n')
that prints
original array
[[ 0 1 2]
[ 3 4 5]
[ 6 7 8]
[ 9 10 11]]
.. row means
[ 1. 4. 7. 10.]
demeaned array
[[-1. 0. 1.]
[-1. 0. 1.]
[-1. 0. 1.]
[-1. 0. 1.]]
.. row means
[0. 0. 0. 0.]
Rule 2 makes it clear why this works as an array with shape (4, 1) can be stretched across the columns for its shape to become (4, 3).
In the third example we will demonstrate that broadcasting can stretch both arrays in the ufunc binary function
a = np.arange(4)
b = np.arange(3)
res = a[:, np.newaxis] + b[np.newaxis, :]
print('result array', res, sep='\n')
that gives
result array
[[0 1 2]
[1 2 3]
[2 3 4]
[3 4 5]]
Strictly speaking the reshaping of the b was not necessary, but it makes things a bit clearer. We can also broadcast the two arrays without applying the ufunc with np.broadcast_arrays or the related and more flexible np.broadcast. For completeness, there are other ways to achieve the same result as with broadcasting, with one example being np.add.outer(a, b) that produces an equal result as
np.array_equal(np.add.outer(a, b), a[:, np.newaxis] + b)
returns True.
Demeaning a high dimensional array across any axis can be generalised with
def demean_axis(arr, axis=1):
means = arr.mean(axis)
indexer = [slice(None)]*arr.ndim
indexer[axis] = np.newaxis
return arr - means[tuple(indexer)]arr = np.linspace(1, 12, 24*3).reshape(6,4,3)
res = demean_axis(arr, axis=1)
We can confirm that np.abs(res.mean(axis=1)).max() is numerically equal to zero. The above function was taken from the book by Wes McKinney but had to be slightly modified to work with the NumPy version used for this article (1.23.4).
As a more real life example, we can use broadcasting to convert a color image to greyscale. The broadcasting part is annotated with comments:
import matplotlib
matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
import io
import numpy as np
from PIL import Image
# read in the original image (png)
with open('landscape_water_lake_nature_trees.png', mode='rb') as f:
image_orig = f.read()
f = io.BytesIO(image_orig)
im = Image.open(f)
image_orig = np.array(im)/255.
print('shape of original image', image_orig.shape, sep='\n')# convert RGBA to RGB (pillow could be used for this)
background = (1., 1., 1.)
row = image_orig.shape[0]
col = image_orig.shape[1]
image_color = np.zeros( (row, col, 3), dtype='float32' )
r, g, b, a = image_orig[:,:,0], image_orig[:,:,1], image_orig[:,:,2], image_orig[:,:,3]
a = np.asarray( a, dtype='float32' )
R, G, B = background
image_color[:,:,0] = r * a + (1.0 - a) * R
image_color[:,:,1] = g * a + (1.0 - a) * G
image_color[:,:,2] = b * a + (1.0 - a) * B
print('shape of image after RGBA to RGB conversion', image_color.shape, sep='\n')
# convert to greyscale
conv = np.array([0.2126, 0.7152, 0.0722])
# --- broadcasting !!! ---
image_grey = (image_color[:,:,:3]*conv).sum(axis=2)
# --- broadcasting !!! ---
print('shape of image after conversion to greyscale', image_grey.shape, sep='\n')
# plot the image
fig = plt.figure(figsize=(8, 4))
axs = fig.subplots(1, 2)
axs[0].axis('off')
axs[0].set_title('RGB image')
axs[0].imshow(image_color)
axs[1].axis('off')
axs[1].set_title('greyscale image')
axs[1].imshow(image_grey, cmap='gray')
axs[0].annotate('',xy=(0.52,0.5),xytext=(0.50,0.5),arrowprops=dict(facecolor='black'),
xycoords='figure fraction', textcoords='figure fraction')
fig.savefig('RGBA_to_greyscale.png')
The code above produces
What is perhaps of interest is that the original is a RGBA image (red green blue alpha), i.e. it also contains a fourth alpha channel to indicate how opaque each pixel is. We converted the RGBA image to RGB by using a white background. This, and pretty much everything in this example, could have been done with Pillow, but we used NumPy as much as possible on purpose (e.g. Pillow can convert to greyscale maintaining the transparency). The last part of the code is some Matplotlib gimmickry to compare the color and greyscale images.
The conversion to greyscale is done using the equation
from this wikipedia page. The part of the code carrying out the multiplication and summation that changes the shape of the image array from (1080, 1920, 3) to (1080, 1920) is commented. You may wonder if it was worth showing such a long example for one line of broadcasting code. This is exactly the point! Broadcasting is concise and without it the code will be much longer and much slower. It is often the case that the magic behind an algorithm is few lines of NumPy operations, often including broadcasting.
Setting values
Most NumPy users associate broadcasting with array addition or multiplication. However, broadcasting applies equally when using indexing to set values. Below you can find a set of examples that we will not comment in detail
# set one row, same value to all columns
a = np.ones((4,3))
a[0] = -1
# array([[-1., -1., -1.],
# [ 1., 1., 1.],
# [ 1., 1., 1.],
# [ 1., 1., 1.]])# set one column, same value to all rows
a = np.ones((4,3))
a[:, 0] = -1
# array([[-1., 1., 1.],
# [-1., 1., 1.],
# [-1., 1., 1.],
# [-1., 1., 1.]])
# set all rows, same value to all elements
a = np.ones((4,3))
a[:] = -1
# array([[-1., -1., -1.],
# [-1., -1., -1.],
# [-1., -1., -1.],
# [-1., -1., -1.]])
# set all rows, different value to each column
a = np.ones((4,3))
b = np.array([-1, -2, -3])
a[:] = b
# array([[-1., -2., -3.],
# [-1., -2., -3.],
# [-1., -2., -3.],
# [-1., -2., -3.]])
# set all rows, different value to each row
a = np.ones((4,3))
b = np.array([-1, -2, -3, -4])
b = b[:, np.newaxis]
a[:] = b
# array([[-1., -1., -1.],
# [-2., -2., -2.],
# [-3., -3., -3.],
# [-4., -4., -4.]])
# set some rows, different value to each column
a = np.ones((4,3))
b = np.array([-1, -2, -3])
a[:3] = b
# array([[-1., -2., -3.],
# [-1., -2., -3.],
# [-1., -2., -3.],
# [ 1., 1., 1.]])
# set some rows, different value to each row
a = np.ones((4,3))
b = np.array([-1, -2, -3])
a[:3] = b[:, np.newaxis]
# array([[-1., -1., -1.],
# [-2., -2., -2.],
# [-3., -3., -3.],
# [ 1., 1., 1.]])
# set some columns, different value to each column
a = np.ones((4,3))
b = np.array([-1, -2])
a[:,:2] = b
# array([[-1., -2., 1.],
# [-1., -2., 1.],
# [-1., -2., 1.],
# [-1., -2., 1.]])
# set some columns, different value to each row
a = np.ones((4,3))
b = np.array([-1, -2, -3, -4])
a[:,:2] = b[:, np.newaxis]
# array([[-1., -1., 1.],
# [-2., -2., 1.],
# [-3., -3., 1.],
# [-4., -4., 1.]])
Broadcasting may seem complex, but it can be easily mastered if a few key principles are kept in mind. The most important principle is that the array shapes are aligned starting from the right. Missing dimensions are filled with ones, always on the left. The two shapes become identical via stretching of dimensions that are equal to one. It may be necessary to do some reshaping with np.newaxis before two (or more!) arrays are broadcastable. Broadcasting is not only used to compute a whole array but also to set some of the values of an array. This is the gist of it. With some practice NumPy broadcasting can lead to surprisingly concise and efficient code. Use it to its full potential!
Definitions, rules and examples
NumPy offers fast calculations via vectorisation that avoids the use of slow Python loops. Vectorisation is also available when using binary ufuncs, such as addition or multiplication, with the added benefit that arrays do not need to have the same shape. Operations with arrays of different shape is known as broadcasting and can be particularly confusing, especially with multidimensional arrays, or when both arrays need to be stretched.
There are many examples and tutorials available, but I find most useful to approach the matter by thinking, and actually memorising, the broadcasting rules. It is then easier to think about any given use case and write the code without relying on trial-and-error.
There are two data analysis and data science books that I would highly recommend. Both of them contain a small section on broadcasting.
Python for Data Analysis by Wes McKinney contains he following broadcasting rule:
Two arrays are compatible for broadcasting if for each trailing dimension (i.e., starting from the end) the axis lengths match or if either of the lengths is 1. Broadcasting is then performed over the missing or length 1 dimensions.
Python Data Science Handbook by Jake VanderPlas contains a more verbose set of broadcasting rules:
Rule 1: If the two arrays differ in their number of dimensions, the shape of the one with fewer dimensions is padded with ones on its leading (left) side.
Rule 2: If the shape of the two arrays does not match in any dimension, the array with shape equal to 1 in that dimension is stretched to match the other shape.
Rule 3: If in any dimension the sizes disagree and neither is equal to 1, an error is raised.
I find the second set of rules easier to follow. The examples below will be using these rules.
Perhaps one of the simplest examples of broadcasting, and a typical pattern, is to subtract the column mean from each column. Following this operation the column means will become numerically equally to zero
a = np.arange(12).reshape(4, 3)
means_columns = a.mean(axis=0)
res = a - means_columns
print('original array', a, sep='\n')
print('.. column means', a.mean(axis=0), sep='\n')
print('demeaned array', res, sep='\n')
print('.. column means', res.mean(axis=0), sep='\n')
that prints
original array
[[ 0 1 2]
[ 3 4 5]
[ 6 7 8]
[ 9 10 11]]
.. column means
[4.5 5.5 6.5]
demeaned array
[[-4.5 -4.5 -4.5]
[-1.5 -1.5 -1.5]
[ 1.5 1.5 1.5]
[ 4.5 4.5 4.5]]
.. column means
[0. 0. 0.]
Let’s see what is going on here. Computing the means of the columns with a.mean(axis=0) produces an one-dimensional array with shape (3,). The two arrays involved in the subtraction differ in shape and hence means_columnsthat has fewer dimensions, will have its shape padded with 1 on the left according to rule 1. Hence, behind the scenes, means_columns will be reshaped to (1, 3). Then according to rule 2 means_columns will be stretched along axis 0 so that its shape becomes (3, 3) to match the shape of a.
In addition to using the rules to predict how the array with fewer dimensions will be stretched we can also use np.broadcast_to that returns a read-only view on the original array with the given shape. The view may be non-contiguous and it is possible that different elements refer to the same memory address
means_columns_bc = np.broadcast_to(means_columns, a.shape)
print(means_columns_bc)
print('base', means_columns_bc.base, sep='\n')
print('strides', means_columns_bc.strides, sep='\n')
that prints
[[4.5 5.5 6.5]
[4.5 5.5 6.5]
[4.5 5.5 6.5]
[4.5 5.5 6.5]]
base
[4.5 5.5 6.5]
strides
(0, 8)
We can see that the base is the original means array (hence it is a view), whilst the stride along the first axis is 0, which means that different elements of the same column refer to the same memory location (see here for a primer on NumPy internals). NumPy is truly optimising memory usage wherever it can!
What if we wanted to demean the rows? The mean of the rows can be computed readily with a.mean(axis=1) that will return a (4,) array. Padding its shape with 1 on the left, means that the array will become (1, 4). According to rule 3 the last dimension of the two arrays disagree and neither is 1. This means that broadcasting will not take place. We could also anticipate this becausenp.broadcast_to(a.mean(axis=1), a.shape) raises an exception informing us that broadcast could not produce the requested shape (4, 3). The incompatibility of shapes can also be seen by executing np.broadcast_shapes(a.shape, a.mean(axis=1).shape) that also raises an exception explaining the shape mismatch. Demeaning the rows can be done by reshaping the row means into a (4, 1) array using either a.mean(axis=1).reshape(-1, 1) or a.mean(axis=1)[:, np.newaxis]
means_rows = a.mean(axis=1)
res = a - means_rows.reshape(-1, 1) # or res = a - means_rows[:, np.newaxis]
print('original array', a, sep='\n')
print('.. row means', a.mean(axis=1), sep='\n')
print('demeaned array', res, sep='\n')
print('.. row means', res.mean(axis=1), sep='\n')
that prints
original array
[[ 0 1 2]
[ 3 4 5]
[ 6 7 8]
[ 9 10 11]]
.. row means
[ 1. 4. 7. 10.]
demeaned array
[[-1. 0. 1.]
[-1. 0. 1.]
[-1. 0. 1.]
[-1. 0. 1.]]
.. row means
[0. 0. 0. 0.]
Rule 2 makes it clear why this works as an array with shape (4, 1) can be stretched across the columns for its shape to become (4, 3).
In the third example we will demonstrate that broadcasting can stretch both arrays in the ufunc binary function
a = np.arange(4)
b = np.arange(3)
res = a[:, np.newaxis] + b[np.newaxis, :]
print('result array', res, sep='\n')
that gives
result array
[[0 1 2]
[1 2 3]
[2 3 4]
[3 4 5]]
Strictly speaking the reshaping of the b was not necessary, but it makes things a bit clearer. We can also broadcast the two arrays without applying the ufunc with np.broadcast_arrays or the related and more flexible np.broadcast. For completeness, there are other ways to achieve the same result as with broadcasting, with one example being np.add.outer(a, b) that produces an equal result as
np.array_equal(np.add.outer(a, b), a[:, np.newaxis] + b)
returns True.
Demeaning a high dimensional array across any axis can be generalised with
def demean_axis(arr, axis=1):
means = arr.mean(axis)
indexer = [slice(None)]*arr.ndim
indexer[axis] = np.newaxis
return arr - means[tuple(indexer)]arr = np.linspace(1, 12, 24*3).reshape(6,4,3)
res = demean_axis(arr, axis=1)
We can confirm that np.abs(res.mean(axis=1)).max() is numerically equal to zero. The above function was taken from the book by Wes McKinney but had to be slightly modified to work with the NumPy version used for this article (1.23.4).
As a more real life example, we can use broadcasting to convert a color image to greyscale. The broadcasting part is annotated with comments:
import matplotlib
matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
import io
import numpy as np
from PIL import Image
# read in the original image (png)
with open('landscape_water_lake_nature_trees.png', mode='rb') as f:
image_orig = f.read()
f = io.BytesIO(image_orig)
im = Image.open(f)
image_orig = np.array(im)/255.
print('shape of original image', image_orig.shape, sep='\n')# convert RGBA to RGB (pillow could be used for this)
background = (1., 1., 1.)
row = image_orig.shape[0]
col = image_orig.shape[1]
image_color = np.zeros( (row, col, 3), dtype='float32' )
r, g, b, a = image_orig[:,:,0], image_orig[:,:,1], image_orig[:,:,2], image_orig[:,:,3]
a = np.asarray( a, dtype='float32' )
R, G, B = background
image_color[:,:,0] = r * a + (1.0 - a) * R
image_color[:,:,1] = g * a + (1.0 - a) * G
image_color[:,:,2] = b * a + (1.0 - a) * B
print('shape of image after RGBA to RGB conversion', image_color.shape, sep='\n')
# convert to greyscale
conv = np.array([0.2126, 0.7152, 0.0722])
# --- broadcasting !!! ---
image_grey = (image_color[:,:,:3]*conv).sum(axis=2)
# --- broadcasting !!! ---
print('shape of image after conversion to greyscale', image_grey.shape, sep='\n')
# plot the image
fig = plt.figure(figsize=(8, 4))
axs = fig.subplots(1, 2)
axs[0].axis('off')
axs[0].set_title('RGB image')
axs[0].imshow(image_color)
axs[1].axis('off')
axs[1].set_title('greyscale image')
axs[1].imshow(image_grey, cmap='gray')
axs[0].annotate('',xy=(0.52,0.5),xytext=(0.50,0.5),arrowprops=dict(facecolor='black'),
xycoords='figure fraction', textcoords='figure fraction')
fig.savefig('RGBA_to_greyscale.png')
The code above produces
What is perhaps of interest is that the original is a RGBA image (red green blue alpha), i.e. it also contains a fourth alpha channel to indicate how opaque each pixel is. We converted the RGBA image to RGB by using a white background. This, and pretty much everything in this example, could have been done with Pillow, but we used NumPy as much as possible on purpose (e.g. Pillow can convert to greyscale maintaining the transparency). The last part of the code is some Matplotlib gimmickry to compare the color and greyscale images.
The conversion to greyscale is done using the equation
from this wikipedia page. The part of the code carrying out the multiplication and summation that changes the shape of the image array from (1080, 1920, 3) to (1080, 1920) is commented. You may wonder if it was worth showing such a long example for one line of broadcasting code. This is exactly the point! Broadcasting is concise and without it the code will be much longer and much slower. It is often the case that the magic behind an algorithm is few lines of NumPy operations, often including broadcasting.
Setting values
Most NumPy users associate broadcasting with array addition or multiplication. However, broadcasting applies equally when using indexing to set values. Below you can find a set of examples that we will not comment in detail
# set one row, same value to all columns
a = np.ones((4,3))
a[0] = -1
# array([[-1., -1., -1.],
# [ 1., 1., 1.],
# [ 1., 1., 1.],
# [ 1., 1., 1.]])# set one column, same value to all rows
a = np.ones((4,3))
a[:, 0] = -1
# array([[-1., 1., 1.],
# [-1., 1., 1.],
# [-1., 1., 1.],
# [-1., 1., 1.]])
# set all rows, same value to all elements
a = np.ones((4,3))
a[:] = -1
# array([[-1., -1., -1.],
# [-1., -1., -1.],
# [-1., -1., -1.],
# [-1., -1., -1.]])
# set all rows, different value to each column
a = np.ones((4,3))
b = np.array([-1, -2, -3])
a[:] = b
# array([[-1., -2., -3.],
# [-1., -2., -3.],
# [-1., -2., -3.],
# [-1., -2., -3.]])
# set all rows, different value to each row
a = np.ones((4,3))
b = np.array([-1, -2, -3, -4])
b = b[:, np.newaxis]
a[:] = b
# array([[-1., -1., -1.],
# [-2., -2., -2.],
# [-3., -3., -3.],
# [-4., -4., -4.]])
# set some rows, different value to each column
a = np.ones((4,3))
b = np.array([-1, -2, -3])
a[:3] = b
# array([[-1., -2., -3.],
# [-1., -2., -3.],
# [-1., -2., -3.],
# [ 1., 1., 1.]])
# set some rows, different value to each row
a = np.ones((4,3))
b = np.array([-1, -2, -3])
a[:3] = b[:, np.newaxis]
# array([[-1., -1., -1.],
# [-2., -2., -2.],
# [-3., -3., -3.],
# [ 1., 1., 1.]])
# set some columns, different value to each column
a = np.ones((4,3))
b = np.array([-1, -2])
a[:,:2] = b
# array([[-1., -2., 1.],
# [-1., -2., 1.],
# [-1., -2., 1.],
# [-1., -2., 1.]])
# set some columns, different value to each row
a = np.ones((4,3))
b = np.array([-1, -2, -3, -4])
a[:,:2] = b[:, np.newaxis]
# array([[-1., -1., 1.],
# [-2., -2., 1.],
# [-3., -3., 1.],
# [-4., -4., 1.]])
Broadcasting may seem complex, but it can be easily mastered if a few key principles are kept in mind. The most important principle is that the array shapes are aligned starting from the right. Missing dimensions are filled with ones, always on the left. The two shapes become identical via stretching of dimensions that are equal to one. It may be necessary to do some reshaping with np.newaxis before two (or more!) arrays are broadcastable. Broadcasting is not only used to compute a whole array but also to set some of the values of an array. This is the gist of it. With some practice NumPy broadcasting can lead to surprisingly concise and efficient code. Use it to its full potential!
Denial of responsibility! Techno Blender is an automatic aggregator of the all world’s media. In each content, the hyperlink to the primary source is specified. All trademarks belong to their rightful owners, all materials to their authors. If you are the owner of the content and do not want us to publish your materials, please contact us by email – [email protected]. The content will be deleted within 24 hours.