NumPy: broadcasting, axes and vectorisation

NumPy is fast because the loop runs in compiled code over contiguous memory instead of in the interpreter over boxed objects. Writing good NumPy means expressing the computation so that loop can happen, and the two things you must be able to do in your head are broadcasting and axes.

Proposition 2.1

The broadcasting rules

Align the shapes from the *trailing* dimension. Two dimensions are compatible if they are equal or one of them is 1; a missing leading dimension is treated as 1. The result takes the larger of each pair, and no data is copied — the size-1 axis is simply reused.

Holds when

(100, 5) with (5,) gives (100, 5) — the row is applied to every row.
(100, 5) with (100, 1) gives (100, 5) — the column is applied to every column.
(100, 5) with (100,) is an error: the trailing dimensions 5 and 100 do not match.

The [:, None] idiom. The third example above is the one that bites, and the fix is to make the intent explicit: a / b[:, None] turns the shape-(100,) vector into (100, 1), which then broadcasts down the columns. Getting into the habit of writing the None rather than relying on a lucky alignment means the code says which axis you meant, and it fails loudly rather than silently doing the wrong thing when a shape changes.

Call on shape `(rows, cols)`	Result shape	Meaning
`a.sum()`	`()`	Everything
`a.sum(axis=0)`	`(cols,)`	Down the rows — one number per column
`a.sum(axis=1)`	`(rows,)`	Across the columns — one number per row
`a.sum(axis=0, keepdims=True)`	`(1, cols)`	Keeps the axis, so it broadcasts back

Table 2.2 · Axes, without guessing. The reliable way to remember it: axis=k is the axis that *disappears*. Summing a (dates, assets) array over axis=0 collapses the dates and leaves one number per asset.

Example 2.3

Standardise a (250, 500) array of returns cross-sectionally — within each date. What do you write?

Show the worked solutionHide the worked solution

Worked solution

Formula $z = \frac{x - \mu_{\text{row}}}{\sigma_{\text{row}}}$
Substitute $\text{dates are axis } 0, \text{ assets are axis } 1$
Solve $\texttt{mu = a.mean(axis=1, keepdims=True)} \Rightarrow (250, 1)$
$\texttt{sd = a.std(axis=1, keepdims=True)} \Rightarrow (250, 1)$
$\texttt{(a - mu) / sd} \Rightarrow (250, 500)$
Answer $\texttt{keepdims=True} \text{ is what makes it broadcast}$

Sanity check. Without keepdims the mean has shape (250,), which broadcasts against the *columns* and standardises by date across the wrong axis — silently, with no error. That is the single most common NumPy bug in research code.

NumPy: broadcasting, axes and vectorisation

After this lesson you should be able to

The broadcasting rules

The rest of this lesson is in Premium