NumPy Vectorisation and Broadcasting Explained

Vectorised code works on whole arrays at once instead of looping over elements in Python. It is shorter to write and much faster, because the loop runs inside NumPy's compiled code.

Element-wise arithmetic#

Operators between two arrays of the same shape work element by element:

import numpy as np

price = np.array([100, 250, 80])
quantity = np.array([3, 1, 5])

print(price * quantity)          # revenue per product
print(price + 10)                # add to every element
print(price / 4)
print(price // 30, price % 30)
print(price ** 2)

[300 250 400]
[110 260  90]
[25.  62.5 20. ]
[3 8 2] [10 10 20]
[10000 62500  6400]

Comparisons give boolean arrays#

import numpy as np

scores = np.array([72, 45, 90, 38, 66])
passed = scores >= 50
print(passed)
print(passed.sum(), "passed")     # True counts as 1
print(np.where(passed, "pass", "fail"))

[ True False  True False  True]
3 passed
['pass' 'fail' 'pass' 'fail' 'pass']

np.where(condition, a, b) picks from a where the condition is true and from b elsewhere: the vectorised version of an if/else.

Universal functions#

NumPy's universal functions (ufuncs) apply a maths function to every element:

import numpy as np

x = np.array([1, 4, 9, 16])
print(np.sqrt(x))
print(np.log10(np.array([1, 10, 100, 1000])))
print(np.round(np.exp(np.array([0, 1, 2])), 3))
print(np.abs(np.array([-3, 2, -1])))
print(np.maximum(np.array([1, 5, 3]), np.array([4, 2, 6])))

[1. 2. 3. 4.]
[0. 1. 2. 3.]
[1.    2.718 7.389]
[3 2 1]
[4 5 6]

Use np.sqrt, not math.sqrt: the math functions only accept single numbers.

Broadcasting#

When shapes differ, NumPy broadcasts: it stretches the smaller array (without copying) so the shapes match. The rule, comparing shapes from the right: two dimensions are compatible if they are equal or one of them is 1.

import numpy as np

sales = np.array([
    [10, 20, 30],      # store A: three products
    [15, 25, 35],      # store B
])
prices = np.array([2.0, 5.0, 1.5])    # one price per product (shape (3,))

print(sales.shape, prices.shape)
print(sales * prices)                  # each row multiplied by the prices

store_bonus = np.array([[100], [200]]) # shape (2, 1): one value per store
print(sales + store_bonus)

(2, 3) (3,)
[[ 20.  100.   45. ]
 [ 30.  125.   52.5]]
[[110 120 130]
 [215 225 235]]

• (2, 3) with (3,): the prices row is reused for each store.
• (2, 3) with (2, 1): each store's bonus is reused across its three products.

A common use is normalising columns: subtract each column's mean from every row with data - data.mean(axis=0).

import numpy as np

a = np.array([[1, 2, 3], [4, 5, 6]])
try:
    a + np.array([1, 2])          # shapes (2, 3) and (2,) do not match from the right
except ValueError as e:
    print("Error:", e)

Error: operands could not be broadcast together with shapes (2,3) (2,)

How much faster?#

Timing depends on the computer, so instead of quoting seconds, compare a loop and a vectorised sum yourself in a notebook with %timeit:

import numpy as np

values = np.arange(1_000_000)

def loop_sum(arr):
    total = 0
    for v in arr:
        total += v
    return total

print(loop_sum(values) == values.sum())

True

Both give the same answer. Run %timeit loop_sum(values) and %timeit values.sum() in Jupyter to compare. On the laptop used to check this course, the loop took about 85 ms and values.sum() about 0.5 ms: roughly 160 times faster. Your numbers will differ, but the gap is always large.

Common mistakes#

Practice#

import numpy as np

scores = np.array([[60, 70, 80], [90, 85, 95], [55, 65, 50]])
weights = np.array([0.2, 0.3, 0.5])

weighted = (scores * weights).sum(axis=1)
print(np.round(weighted, 1))
print(np.where(weighted >= 65, "pass", "fail"))

[73.  91.  55.5]
['pass' 'pass' 'fail']

Key takeaway#

Next: Indexing, Slicing and Boolean Masks.