Skip to main content

Module 4 — NumPy: arrays, broadcasting and vectorized computation

NumPy is the foundation of the whole ecosystem: pandas, scikit-learn, and even PyTorch tensors borrow its concepts. Its single object — the homogeneous ndarray — and its single principle — vectorize instead of looping — are enough to speed computations up by a factor of 10 to 100.

The ndarray: a homogeneous, typed block

Unlike a Python list (references to scattered objects), a NumPy array is a contiguous memory block of same-type elements. That homogeneity is what buys the speed: operations drop down to compiled C instead of interpreting element by element.

import numpy as np

a = np.array([1.5, 2.0, 3.5])
m = np.zeros((3, 4)) # 3×4 matrix of zeros
x = np.arange(0, 10, 0.5) # 0, 0.5, 1.0, …, 9.5
g = np.random.default_rng(42).normal(size=(1000, 3)) # reproducible randomness

a.shape # (3,) — the dimensions
a.dtype # float64 — the single element type
m.ndim # 2 — the number of axes

shape and dtype are the two attributes to check first whenever anything behaves strangely: half of NumPy bugs are unexpected shapes.

Vectorization: the operation applies to the whole array

# Pure Python style — slow
results = []
for price in price_list:
results.append(price * 1.15)

# NumPy style — 10 to 100 times faster, and more readable
results = prices * 1.15

All arithmetic operations and the universal functions (np.log, np.exp, np.sqrt…) apply element-wise to the entire array. The list comprehension from module 2 becomes a loop-free expression:

standardized = (x - x.mean()) / x.std()      # standardization in one line

The professional rule: if you are writing a for loop over a NumPy array, stop — there is almost certainly a vectorized formulation, faster and closer to the mathematics.

Broadcasting: operating between different shapes

NumPy automatically stretches arrays of compatible shapes:

m = np.array([[1., 2., 3.],
[4., 5., 6.]]) # shape (2, 3)

m * 10 # scalar broadcast everywhere
m - m.mean(axis=0) # subtracts EACH COLUMN's mean — shape (3,) broadcast over (2, 3)

The formal rule: two dimensions are compatible if they are equal or one of them is 1, aligning shapes from the right. In practice, the case that matters is the one above: centering or scaling columns without a loop. That is exactly the standardization step that precedes most models.

When broadcasting goes wrong silently

Subtracting a shape-(n,) array from a shape-(n, 1) array produces an (n, n) matrix — with no error. If a computation returns an absurd shape, look for unintended broadcasting; x.reshape(-1, 1) and x.ravel() straighten shapes out.

Boolean masks: filtering by condition

The most used construct in all of NumPy — and the exact mechanism behind the pandas filters of the next module:

ages = np.array([22, 35, 58, 41, 17, 63])

mask = ages >= 40 # array([False, False, True, True, False, True])
ages[mask] # array([58, 41, 63]) — selection
(ages >= 40).mean() # 0.5 — proportion (True counts as 1)

# Combined conditions: & (and), | (or), with MANDATORY parentheses
active_seniors = ages[(ages >= 40) & (ages < 65)]

# Conditional replacement
capped = np.where(ages > 60, 60, ages)

The parentheses around each condition are not optional: & and | bind tighter than comparisons, and forgetting them produces a cryptic error.

Aggregations and axes: summarizing in the right direction

grades = np.array([[12, 15, 9],
[14, 11, 16]]) # 2 students × 3 subjects

grades.mean() # 12.83 — global mean
grades.mean(axis=0) # [13. 13. 12.5] — per subject (rows collapse)
grades.mean(axis=1) # [12. 13.67] — per student (columns collapse)

The mnemonic that sticks: axis names the axis that disappears. axis=0 collapses the rows, leaving one value per column. The same aggregations exist everywhere: sum, min, max, std, argmax (the index of the maximum — the one that yields a neural network's predicted class).

Two real-world subtleties: missing values propagate (np.nan contaminates any sum — use np.nanmean and its family), and NumPy slices are views over the same memory, not copies (.copy() to decouple).

Key takeaways

  • The ndarray is homogeneous and typed; shape and dtype are the first diagnostic reflexes.
  • Vectorizing replaces loops: element-wise operations, 10 to 100 times faster, more readable.
  • Broadcasting lets you operate between compatible shapes — centering columns in one line — but can silently manufacture absurd shapes.
  • Boolean masks: filter, count, replace by condition; mandatory parentheses with & and |.
  • axis = the axis that disappears; np.nanmean for missing values; slices are views.

Next module: pandas, which puts labels — column names, an index — on these arrays and becomes the daily tool for tabular data work.