NumPy Tutorial: Your First Steps Into Data Science in Python – Real Python


So far, you’ve seen a couple of smaller examples of broadcasting, but the topic will start to make more sense the more examples you see. Fundamentally, it functions around one rule: arrays can be broadcast against each other if their dimensions match or if one of the arrays has a size of 1.

If the arrays match in size along an axis, then elements will be operated on element-by-element, similar to how the built-in Python function zip() works. If one of the arrays has a size of 1 in an axis, then that value will be broadcast along that axis, or duplicated as many times as necessary to match the number of elements along that axis in the other array.

Here’s a quick example. Array A has the shape (4, 1, 8), and array B has the shape (1, 6, 8). Based on the rules above, you can operate on these arrays together:

All three axes successfully follow the rule.

>>>
In [1]: import numpy as np In [2]: A = np.arange(32).reshape(4, 1, 8) In [3]: A
Out[3]:
array([[[ 0, 1, 2, 3, 4, 5, 6, 7]],  [[ 8, 9, 10, 11, 12, 13, 14, 15]],  [[16, 17, 18, 19, 20, 21, 22, 23]],  [[24, 25, 26, 27, 28, 29, 30, 31]]]) In [4]: B = np.arange(48).reshape(1, 6, 8) In [5]: B
Out[5]:
array([[[ 0, 1, 2, 3, 4, 5, 6, 7],
 [ 8, 9, 10, 11, 12, 13, 14, 15],
 [16, 17, 18, 19, 20, 21, 22, 23],
 [24, 25, 26, 27, 28, 29, 30, 31],
 [32, 33, 34, 35, 36, 37, 38, 39],
 [40, 41, 42, 43, 44, 45, 46, 47]]])

A has 4 planes, each with 1 row and 8 columns. B has only 1 plane with 6 rows and 8 columns. Watch what NumPy does for you when you try to do a calculation between them!

>>>
In [7]: A + B
Out[7]:
array([[[ 0, 2, 4, 6, 8, 10, 12, 14],
 [ 8, 10, 12, 14, 16, 18, 20, 22],
 [16, 18, 20, 22, 24, 26, 28, 30],
 [24, 26, 28, 30, 32, 34, 36, 38],
 [32, 34, 36, 38, 40, 42, 44, 46],
 [40, 42, 44, 46, 48, 50, 52, 54]],  [[ 8, 10, 12, 14, 16, 18, 20, 22],
 [16, 18, 20, 22, 24, 26, 28, 30],
 [24, 26, 28, 30, 32, 34, 36, 38],
 [32, 34, 36, 38, 40, 42, 44, 46],
 [40, 42, 44, 46, 48, 50, 52, 54],
 [48, 50, 52, 54, 56, 58, 60, 62]],  [[16, 18, 20, 22, 24, 26, 28, 30],
 [24, 26, 28, 30, 32, 34, 36, 38],
 [32, 34, 36, 38, 40, 42, 44, 46],
 [40, 42, 44, 46, 48, 50, 52, 54],
 [48, 50, 52, 54, 56, 58, 60, 62],
 [56, 58, 60, 62, 64, 66, 68, 70]],  [[24, 26, 28, 30, 32, 34, 36, 38],
 [32, 34, 36, 38, 40, 42, 44, 46],
 [40, 42, 44, 46, 48, 50, 52, 54],
 [48, 50, 52, 54, 56, 58, 60, 62],
 [56, 58, 60, 62, 64, 66, 68, 70],
 [64, 66, 68, 70, 72, 74, 76, 78]]])

The way broadcasting works is that NumPy duplicates the plane in B three times so that you have a total of four, matching the number of planes in A. It also duplicates the single row in A five times for a total of six, matching the number of rows in B. Then it adds each element in the newly expanded A array to its counterpart in the same location in B. The result of each calculation shows up in the corresponding location of the output.

Once again, even though you can use words like “plane,” “row,” and “column” to describe how the shapes in this example are broadcast to create matching three-dimensional shapes, things get more complicated at higher dimensions. A lot of times, you’ll have to simply follow the broadcasting rules and do lots of print-outs to make sure things are working as planned.

Understanding broadcasting is an important part of mastering vectorized calculations, and vectorized calculations are the way to write clean, idiomatic NumPy code.