NumPy — Arrays, Vectors, and Why Loops Are Slow

NumPy — Arrays, Vectors, and Why Loops Are Slow

NumPy is the foundation of the entire Python data and AI stack. pandas DataFrames are built on NumPy arrays; scikit-learn takes them as input; embeddings (Module 4) are NumPy vectors. You don't need to be a NumPy expert, but you need its core idea: operate on whole arrays at once, not element by element.

Arrays

import numpy as np

durations = np.array([412, 2310, 188, 95, 1980])   # ms
durations.dtype     # dtype('int64')
durations.shape     # (5,)
durations.mean()    # 997.0
durations.max(), durations.argmax()   # (2310, 1) — value and its index

Unlike a Python list, an array holds one type packed tightly in memory. That's what makes it fast.

Vectorisation: No Loops

Arithmetic applies to every element at once:

seconds = durations / 1000
over_budget = durations > 1000          # array([False, True, False, False, True])
durations[over_budget]                  # array([2310, 1980]) — boolean indexing
over_budget.sum()                       # 2 — True counts as 1
np.percentile(durations, 95)            # p95 latency

The equivalent Python loop does the same work one element at a time, in the interpreter. NumPy does it in optimised C. On a million latency samples, the difference is typically 50–100×:

samples = np.random.default_rng(0).normal(300, 50, 1_000_000)

# Python loop: ~150 ms
total = 0.0
for s in samples:
    total += s

# NumPy: ~1 ms
total = samples.sum()

That's why idiomatic AI code almost never loops over individual numbers.

Shapes and 2-D Arrays

Data usually comes as a matrix: rows are examples, columns are features.

# 4 test runs × 3 metrics (duration_ms, retries, assertions)
runs = np.array([
    [412, 0, 12],
    [2310, 2, 30],
    [188, 0, 4],
    [1980, 1, 28],
])
runs.shape            # (4, 3)
runs[:, 0]            # first column — all durations
runs[1]               # second row
runs.mean(axis=0)     # mean of each column
runs.mean(axis=1)     # mean of each row

axis=0 collapses rows (giving one value per column); axis=1 collapses columns. Getting the axis wrong is one of the most common NumPy bugs — check .shape whenever results look odd.

Broadcasting

NumPy can combine arrays of different shapes when they're compatible. The common case: apply a per-column operation to every row.

col_means = runs.mean(axis=0)          # shape (3,)
col_stds = runs.std(axis=0)            # shape (3,)
standardised = (runs - col_means) / col_stds   # (4,3) op (3,) → each row

That one line standardises every feature — a typical step before training many ML models.

Vectors and Similarity — Preview of Embeddings

In Module 4, text is turned into vectors of a few hundred numbers. "How similar are these two texts?" becomes "how close are these two vectors?" — measured with cosine similarity:

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))

a = np.array([0.2, 0.9, 0.1])
b = np.array([0.25, 0.85, 0.05])
cosine_similarity(a, b)   # ≈ 0.99 — very similar direction

@ is matrix/dot product. And it vectorises too — comparing one query vector against 10,000 document vectors is a single matrix multiplication:

scores = docs @ query / (np.linalg.norm(docs, axis=1) * np.linalg.norm(query))
top5 = np.argsort(scores)[::-1][:5]

Random Numbers, Reproducibly

ML is full of randomness — shuffling, splitting, initialising. Always use a seeded generator so results are reproducible:

rng = np.random.default_rng(seed=42)
rng.integers(0, 100, size=5)

A test or experiment that can't be reproduced can't be debugged.

Testing Numeric Code

Never compare floats with ==. Use tolerances:

np.testing.assert_allclose(result, expected, rtol=1e-6)

np.testing gives clear failure messages showing which elements differ — far better than assert (a == b).all().