NumPy: thinking in arrays

Python’s entire scientific ecosystem rests on NumPy: pandas stores its columns in NumPy arrays, scikit-learn expects NumPy arrays, and Keras tensors are their direct cousins. Before any model, you need to learn to think in arrays rather than in loops.

Why not plain Python lists?

A Python list holds arbitrary objects, scattered across memory. A NumPy array (ndarray) holds values of a single type, stored side by side. As a result, operations run in C, on contiguous memory blocks, often 50 to 100 times faster than a Python loop.

import numpy as np

prices = np.array([12.5, 8.0, 23.9, 4.2])
vat = prices * 0.18      # one operation, applied to the whole array
print(vat)               # [2.25  1.44  4.302 0.756]

No loop: this is vectorisation, the number one reflex to build.

Shape, dimensions, type

Three attributes to always keep in mind:

x = np.zeros((3, 4))     # 3 rows, 4 columns
x.shape                  # (3, 4)
x.ndim                   # 2
x.dtype                  # float64

In machine learning, the convention is almost always the same: one row = one example, one column = one feature. A dataset of 1,000 houses described by 5 variables is an array of shape (1000, 5).

To change the shape without copying the data:

v = np.arange(12)        # [0, 1, ..., 11]
m = v.reshape(3, 4)      # 3 x 4
m.T.shape                # (4, 3): the transpose

Indexing and masks

m[0]          # first row
m[:, 1]       # second column
m[1:, :2]     # rows 1 and 2, columns 0 and 1

grades = np.array([12, 7, 15, 9, 18])
grades[grades >= 10]        # [12 15 18]: filtering with a boolean mask
grades[grades < 10] = 10    # conditional replacement, no loop

Boolean masks show up everywhere: filtering data, counting errors, selecting classes.

Broadcasting

NumPy can combine arrays of different shapes when it makes sense. A classic example: centring each column of a dataset.

X = np.random.default_rng(0).normal(size=(100, 3))
X_centred = X - X.mean(axis=0)   # (100, 3) - (3,) -> the mean is "broadcast"
X_centred.mean(axis=0).round(10) # [0. 0. 0.]

axis=0 aggregates along the rows (one result per column), axis=1 along the columns (one result per row). Getting the axis wrong is the most common mistake early on: always check the shape of the result.

Linear algebra, where it all happens

A linear model is a matrix product. With $X$ of shape $(n, d)$ and weights $w$ of shape $(d,)$ :

$$\hat{y} = Xw + b$$

rng = np.random.default_rng(42)
X = rng.normal(size=(200, 2))
w_true = np.array([3.0, -2.0])
y = X @ w_true + 1.0 + rng.normal(scale=0.1, size=200)

# Least squares: add a column of ones for the bias
Xb = np.c_[X, np.ones(len(X))]
w, *_ = np.linalg.lstsq(Xb, y, rcond=None)
print(w.round(2))   # ≈ [ 3.  -2.   1. ]

This is exactly the linear regression from my dedicated article, written in three lines.

Exercises to validate the step

1. Normalise an (n, d) array so that each column has mean 0 and standard deviation 1, without a loop.
2. Compute the Euclidean distance matrix between all points of an (n, 2) cloud using broadcasting.
3. Reimplement linear regression gradient descent in pure NumPy.
4. Load an image as a (height, width, 3) array and convert it to greyscale.

Once these four exercises come without looking up the syntax at every line, the step is validated.

Resources

• NumPy: the absolute basics for beginners, the official documentation — excellent.
• The “NumPy” and “Intro to Machine Learning” exercises on Kaggle Learn.