Chapter 3: Vectorize Everything

A Python loop over a million numbers is slow, because every pass goes back through the interpreter. NumPy’s answer is vectorization: write one expression over whole arrays and let NumPy run the loop in compiled code. Functions like np.sqrt, and the ones behind + and * on arrays, are ufuncs: they work value by value on any shape. A result that is a single number comes back as a NumPy scalar, which the prompt shows with its type, like np.int64(…).

import numpy as np
cents = np.array([1999, 549, 375])
qty = np.array([2, 10, 4])
total = 0
for price, count in zip(cents, qty):
    total += price * count
total

np.int64(10988)

(cents * qty).sum()

np.int64(10988)

np.sqrt(np.array([1.0, 4.0, 9.0]))

array([1., 2., 3.])

Broadcasting is how arrays of different shapes still combine. NumPy lines the shapes up from the right. Two lengths fit if they are equal or one of them is 1, and an axis of length 1, or a missing one, is stretched to match.

A (3, 4) array plus a (4,) row adds the row to every row. A (3, 4) plus a (3,) fails, and the error names both shapes.

import numpy as np
a = np.arange(12).reshape(3, 4)
row = np.array([10, 20, 30, 40])
a + row

array([[10, 21, 32, 43],
       [14, 25, 36, 47],
       [18, 29, 40, 51]])

col = np.array([[100], [200], [300]])
(a + col).shape

(3, 4)

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

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

ML preprocessing uses this line again and again: normalizing a batch so each column has mean 0 and spread 1.

import numpy as npX = np.array([[1, 2, 30], [3, 4, 10], [5, 6, 20], [7, 8, 40]])mu = X.mean(axis=0)X - musd = X.std(axis=0)Z = (X - mu) / sdZ.mean(axis=0).round(6), Z.std(axis=0)

A reduction turns many values into fewer: sum, mean, max, and argmax, which gives the position of the largest value instead of the value. With no axis, they return one number. With axis= they collapse only that axis, so axis=0 gives one value per column.

keepdims=True keeps the collapsed axis with length 1, so the result still lines up with the original. This is where most people slip. To normalize rows instead of columns, the row means need keepdims: without it they have shape (rows,), which lines up with the columns, so the subtraction fails or, on a square array, subtracts the wrong numbers without an error.

import numpy as np
X = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
X.sum(), X.sum(axis=0), X.sum(axis=1)

(np.float64(21.0), array([5., 7., 9.]), array([ 6., 15.]))

X.argmax(axis=1)

array([2, 2])

X.mean(axis=1).shape, X.mean(axis=1, keepdims=True).shape

((2,), (2, 1))

X - X.mean(axis=1, keepdims=True)

array([[-1.,  0.,  1.],
       [-1.,  0.,  1.]])

X - X.mean(axis=1)

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

@ multiplies an (n, k) array by a (k, m) one into (n, m): the inner lengths must match, and they disappear. np.dot gives the same for 2-D arrays and the plain dot product for 1-D ones, but NumPy’s docs prefer @ for 2-D. np.concatenate joins arrays along an axis they already have, and np.stack adds a new one. np.where picks from two arrays by a condition, and np.clip caps values into a range.

import numpy as np
A = np.ones((2, 3))
B = np.ones((3, 4))
(A @ B).shape

(2, 4)

np.dot(np.array([1, 2, 3]), np.array([4, 5, 6]))

np.int64(32)

a = np.zeros((2, 3))
b = np.ones((2, 3))
np.concatenate([a, b]).shape, np.stack([a, b]).shape

((4, 3), (2, 2, 3))

x = np.array([-2, -1, 0, 1, 2])
np.where(x > 0, x, 0)

array([0, 0, 0, 1, 2])

np.clip(x, -1, 1)

array([-1, -1,  0,  1,  1])

One more trap raises no error at all. A slice is a view: it shares memory with the array it came from, so writing into the slice writes into the original. A function that edits a slice of its argument changes the caller’s data, and nothing warns you. Call .copy() when you need a copy of your own, and ask np.shares_memory when you aren’t sure.

import numpy as np
a = np.arange(6)
b = a[:3]
b[0] = 99
a

array([99,  1,  2,  3,  4,  5])

np.shares_memory(a, b)

True

c = a[:3].copy()
c[0] = -1
a

array([99,  1,  2,  3,  4,  5])

To keep an array between runs, np.save writes it to a .npy file and np.load reads it back with its dtype and shape.

import numpy as np
X = np.arange(6, dtype=np.float32).reshape(2, 3)
np.save("batch.npy", X)
Y = np.load("batch.npy")
Y.dtype, Y.shape, np.array_equal(X, Y)

(dtype('float32'), (2, 3), True)

The habit from the last chapter gains two checks. Before you write a loop over an array, look for the one expression that does the same work, and say the shape it will make. Before you write into an array a function was given, ask whether it is a view of someone else’s data.

Homework

The packet has three parts: predict the shapes that broadcasting and reductions produce, write five small functions as whole-array expressions with no loop, then fix two bugs: a slow loop (the test times it) and a function that silently changes its input through a view.