Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

NumPy

🎯 Goal

After reading this chapter:

  • You’ll understand the difference between NumPy ndarray and Python list and know when to use which
  • You’ll learn to write fast code without loops via vectorized operations
  • You’ll be able to work with arrays of different sizes using broadcasting
  • You’ll know the NumPy patterns that appear in 90% of ML code

What to learn

  • Creating ndarray: np.array, np.zeros, np.ones, np.arange, np.linspace, np.random
  • Array attributes: shape, dtype, ndim, size
  • Indexing and slicing (1-D, 2-D, boolean, fancy)
  • Reshape, transpose, concatenate, stack, split
  • Arithmetic operations and broadcasting
  • Universal functions (ufuncs): np.sin, np.exp, np.log, etc.
  • Aggregations: sum, mean, max, min, argmax, axis parameter
  • Linear algebra (np.linalg)
  • Random sampling (np.random)

Libraries

pip install numpy

NumPy version 1.26+ or 2.x is recommended.

Important topics

Why is NumPy faster than Python list?

# Python list — each element is a separate PyObject (slow)
py_list = [1, 2, 3, 1_000_000]
# NumPy — a single C array block (fast)
np_arr = np.array([1, 2, 3, 1_000_000], dtype=np.int64)

NumPy:

  • Written in C, uses SIMD instructions
  • Single dtype (e.g., all int64) — cache-friendly
  • Vectorized: arr * 2 — one operation on the whole array

Bench: multiplying 1M elements by 2 — list ~50ms, NumPy ~1ms (50x faster).

Broadcasting

NumPy’s most powerful feature. Working with arrays of different sizes:

# Adding (3,) vector to (3, 3) matrix
A = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])  # shape (3, 3)
b = np.array([10, 20, 30])                        # shape (3,)
result = A + b
# [[11, 22, 33], [14, 25, 36], [17, 28, 39]]

NumPy automatically “broadcasts” b to each row.

**Rule:**dimensions are compared from the end. They must be equal, or one of them must be 1.

Axis concept

For a 2-D array:

  • axis=0 — row (vertical, “down the rows”)
  • axis=1 — column (horizontal, “across the columns”)
A = np.array([[1, 2], [3, 4], [5, 6]])  # shape (3, 2)
A.sum(axis=0)  # [9, 12] — sum of each column
A.sum(axis=1)  # [3, 7, 11] — sum of each row

Code examples

Array creation and basic operations

import numpy as np

# Creation methods
a = np.array([1, 2, 3, 4])
zeros = np.zeros((3, 4))            # 3×4 zeros
ones = np.ones((2, 2))               # 2×2 ones
rng = np.arange(0, 10, 2)            # [0, 2, 4, 6, 8]
lin = np.linspace(0, 1, 5)           # 5 evenly distributed numbers [0, 0.25, 0.5, 0.75, 1]
random_arr = np.random.rand(3, 3)    # 3×3 random [0, 1)

# Attributes
print(a.shape, a.dtype, a.ndim, a.size)  # (4,) int64 1 4

Indexing and boolean filtering

arr = np.array([10, 20, 30, 40, 50, 60])

# Slicing
print(arr[1:4])      # [20, 30, 40]
print(arr[::-1])     # reverse

# Boolean indexing — used a lot in ML
mask = arr > 30
print(arr[mask])     # [40, 50, 60]

# Filter and modify simultaneously
arr[arr < 30] = 0
print(arr)           # [0, 0, 30, 40, 50, 60]

# 2-D indexing
M = np.arange(12).reshape(3, 4)
print(M[1, 2])       # 6
print(M[:, 1])       # 2nd column
print(M[1:, :2])     # 1st row, columns 0 and 1

Vectorized operations (no loops)

# With loop (SLOW — don't do this)
arr = np.arange(1_000_000)
result = []
for x in arr:
    result.append(x ** 2 + 3 * x - 5)

# Vectorized (FAST — always do this)
arr = np.arange(1_000_000)
result = arr ** 2 + 3 * arr - 5

# Conditional vectorization
prices = np.array([100, 50, 200, 75, 300])
discounted = np.where(prices > 100, prices * 0.9, prices)
# [100, 50, 180, 75, 270]

Backend integration

NumPy is handy in backend in these places:

1. Fast JSON aggregation

from fastapi import FastAPI
import numpy as np

app = FastAPI()

@app.post("/metrics/")
def process_metrics(values: list[float]):
    arr = np.array(values, dtype=np.float64)
    # In pure Python ~500ms for 1M elements, in NumPy ~5ms
    return {
        "sum": float(arr.sum()),
        "mean": float(arr.mean()),
        "p50": float(np.percentile(arr, 50)),
        "p95": float(np.percentile(arr, 95)),
        "p99": float(np.percentile(arr, 99)),
    }

2. Image processing endpoint

import numpy as np
from PIL import Image
from io import BytesIO

@app.post("/image/grayscale/")
async def to_grayscale(file: UploadFile):
    img = Image.open(file.file)
    arr = np.array(img)
    # RGB → grayscale: luminance formula
    gray = (0.299 * arr[..., 0] + 0.587 * arr[..., 1] + 0.114 * arr[..., 2]).astype(np.uint8)
    out = Image.fromarray(gray)
    buf = BytesIO()
    out.save(buf, format="PNG")
    return Response(content=buf.getvalue(), media_type="image/png")

3. Embedding similarity (for RAG)

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

# Batch search — compare 1000 embeddings with query
def find_similar(query: np.ndarray, embeddings: np.ndarray, top_k: int = 5):
    # embeddings shape: (1000, 384), query shape: (384,)
    sims = embeddings @ query / (np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query))
    top_indices = np.argsort(sims)[::-1][:top_k]
    return top_indices, sims[top_indices]

Resources

🏋️ Exercises

🟢 Easy

  1. Create an array equal to np.arange(0, 100, 5) and filter out odd numbers.
  2. Create a 3x3 random matrix, find the largest element and its index (np.argmax).
  3. Concatenate two arrays [1, 2, 3, 4] and [5, 6, 7, 8] vertically and horizontally (np.vstack, np.hstack).

🟡 Medium

  1. Generate 10000 random numbers and compute x^2 + 2x + 1 in pure Python for loop and NumPy vectorized. Compare with timeit.
  2. Create a (50, 50) random matrix and simulate a chess board (np.indices + broadcasting).
  3. Normalize: bring each column of a random matrix to the 0..1 range (min-max normalization).

🔴 Hard

  1. Cosine similarity API: create a FastAPI endpoint. User sends query: list[float] and database: list[list[float]]. Return Top-K most similar vectors. All NumPy vectorized (no loops).
  2. Sliding window: for 1-D array, compute memory-efficient rolling mean with window_size=k using np.lib.stride_tricks.

Capstone

In notebooks/month-01/01_numpy_basics.ipynb:

  • Simulate a 30-day activity matrix for 1000 users: shape (1000, 30)
  • Calculate each user’s weekly average activity (shape (1000, 4))
  • Find the top 10 most active users
  • Visualize the activity matrix as a heatmap (with Matplotlib)

✅ Checklist

  • I understand the difference between ndarray and Python list
  • Concepts shape, dtype, axis are clear
  • I use boolean indexing, don’t write for loops
  • I know the broadcasting rule and can apply it on small examples
  • I know matrix operations through np.linalg (dot product, inverse, eigvals)
  • I’ve measured how many times my vectorized code is faster than Python loop

Time to move on to Pandas.