🔴 ML / AI

Python + NumPy: Arrays & Vectorisation

📖 Lesson 43 ⏱ 50 min 🧪 5 questions 💻 3 exercises

🎯 Learning Objectives

Create and inspect NumPy arrays: shape, dtype, ndim, size
Use array creation routines: zeros, ones, arange, linspace, random
Index, slice, and reshape arrays including boolean and fancy indexing
Apply vectorised arithmetic and universal functions (ufuncs) — no Python loops
Understand broadcasting rules and apply them to arrays of different shapes
Perform linear algebra operations: dot product, matrix multiply, inverse, eigenvalues
Profile NumPy code and understand when it beats pure Python

1 — Why NumPy?

Pure Python lists are ill-suited for numerical work:

Boxed objects — each element is a heap-allocated Python object with its own type tag and reference count.
No contiguous memory — list elements are pointers scattered across the heap, killing CPU cache performance.
No SIMD — the interpreter cannot auto-vectorise a Python for loop.

NumPy solves all three problems by storing data in a contiguous C array with a single fixed dtype — exactly the same memory layout used by C and Fortran. Operations are delegated to optimised C, BLAS, and LAPACK routines that are often 100× faster than equivalent Python loops.

The benchmark below makes this concrete:

import numpy as np
import time

N = 1_000_000
py_list  = list(range(N))
np_array = np.arange(N, dtype=np.float64)

# Python loop
t0 = time.perf_counter()
result = [x * 2 for x in py_list]
print(f"Python list: {time.perf_counter() - t0:.4f}s")

# NumPy vectorised
t0 = time.perf_counter()
result = np_array * 2
print(f"NumPy array: {time.perf_counter() - t0:.4f}s")
# NumPy is typically 50-200x faster

speed_comparison.py

Install NumPy (if not already present):

pip install numpy

shell

💡 The Golden Rule
Never loop over a NumPy array in Python. If you find yourself writing for x in arr:, there is almost certainly a vectorised operation that does the same thing orders of magnitude faster.

2 — Creating Arrays

import numpy as np

# ── From Python data ──
a = np.array([1, 2, 3, 4, 5])               # 1-D, dtype inferred (int64)
b = np.array([[1, 2, 3], [4, 5, 6]])         # 2-D (2×3 matrix)
c = np.array([1.0, 2, 3])                    # float64 (one float coerces all)
d = np.array([1, 2, 3], dtype=np.float32)   # explicit dtype

# ── Creation routines ──
np.zeros((3, 4))              # 3×4 matrix of 0.0
np.ones((2, 3, 4))            # 2×3×4 tensor of 1.0
np.full((3, 3), 7)            # filled with 7
np.eye(4)                     # 4×4 identity matrix
np.arange(0, 10, 2)           # [0, 2, 4, 6, 8]  like range()
np.linspace(0, 1, 5)          # [0, 0.25, 0.5, 0.75, 1.0]  evenly spaced
np.logspace(0, 3, 4)          # [1, 10, 100, 1000]  log-spaced

# ── Random arrays ──
rng = np.random.default_rng(seed=42)  # recommended: Generator API
rng.random((3, 3))            # uniform [0, 1)
rng.integers(0, 10, size=5)   # random ints in [0, 10)
rng.normal(0, 1, size=(4, 4)) # standard normal
rng.choice([10, 20, 30], size=6, replace=True)

creating_arrays.py

Every array exposes a set of metadata attributes:

arr = np.zeros((3, 4, 5))
print(arr.shape)    # (3, 4, 5)
print(arr.ndim)     # 3
print(arr.size)     # 60
print(arr.dtype)    # float64
print(arr.itemsize) # 8 bytes per element
print(arr.nbytes)   # 480 bytes total

array_attributes.py

Common dtypes at a glance:

Category	dtypes	Notes
Signed integers	`int8`, `int16`, `int32`, `int64`	Default is `int64` on 64-bit platforms
Unsigned integers	`uint8`, `uint16`, `uint32`, `uint64`	`uint8` common for image pixel data
Floating point	`float16`, `float32`, `float64`	Default is `float64`; ML models often use `float32`
Complex	`complex64`, `complex128`	Pairs of float32 / float64
Boolean	`bool`	Stored as 1 byte per element
Object	`object`	Holds arbitrary Python objects — avoids NumPy optimisations

3 — Indexing, Slicing & Reshaping

import numpy as np

a = np.arange(12).reshape(3, 4)
# array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11]])

# ── Basic indexing ──
a[0]        # first row: [0, 1, 2, 3]
a[-1]       # last row:  [8, 9, 10, 11]
a[1, 2]     # row 1, col 2 → 6
a[0, -1]    # row 0, last col → 3

# ── Slicing ──
a[:, 1]     # all rows, col 1 → [1, 5, 9]
a[1:, 2:]   # rows 1+, cols 2+ → [[6,7],[10,11]]
a[::2]      # every other row → [[0,1,2,3],[8,9,10,11]]

# ── Boolean indexing ──
mask = a > 6
a[mask]           # [7, 8, 9, 10, 11]
a[a % 2 == 0]     # all even elements

# ── Fancy indexing ──
rows = np.array([0, 2])
cols = np.array([1, 3])
a[rows, cols]     # [a[0,1], a[2,3]] = [1, 11]

# ── Reshaping ──
a.reshape(4, 3)     # new shape (must match total elements)
a.reshape(-1)       # flatten to 1-D (same as a.ravel())
a.reshape(2, -1)    # 2 rows, infer columns → (2, 6)
a.T                 # transpose
a.flatten()         # copy as 1-D (vs ravel which may return a view)
np.newaxis          # add dimension: a[np.newaxis, :] → (1, 3, 4)

indexing.py

⚠️ Slices Are Views
NumPy slices return views of the original array — modifying a slice modifies the original data. Use .copy() when you need an independent array: sub = a[1:, 2:].copy()

4 — Vectorised Operations & ufuncs

Every standard arithmetic operator is overloaded to work element-wise on NumPy arrays. Universal functions (ufuncs) are compiled C functions that accept arrays and broadcast automatically.

import numpy as np

a = np.array([1.0, 4.0, 9.0, 16.0])
b = np.array([2.0, 2.0, 3.0,  4.0])

# ── Element-wise arithmetic (all vectorised) ──
a + b          # [3, 6, 12, 20]
a - b          # [-1, 2, 6, 12]
a * b          # [2, 8, 27, 64]
a / b          # [0.5, 2, 3, 4]
a ** 2         # [1, 16, 81, 256]
a % 3          # [1, 1, 0, 1]

# ── Universal functions (ufuncs) ──
np.sqrt(a)                            # [1, 2, 3, 4]
np.exp(np.array([0, 1, 2]))           # [1, e, e²]
np.log(a)                             # natural log
np.log2(a)
np.log10(a)
np.sin(np.linspace(0, np.pi, 5))
np.abs(np.array([-3, -1, 2]))         # [3, 1, 2]
np.maximum(a, b)                      # element-wise max → [2, 4, 9, 16]
np.where(a > 5, a, 0)                 # conditional: keep a where a>5, else 0

# ── Aggregate functions ──
a.sum()              # 30.0
a.mean()             # 7.5
a.std()              # standard deviation
a.var()              # variance
a.min(), a.max()     # 1.0, 16.0
a.argmin(), a.argmax()  # index of min/max → 0, 3
np.cumsum(a)         # cumulative sum
np.prod(a)           # product

# ── Axis argument for 2-D ──
m = np.array([[1, 2, 3], [4, 5, 6]])
m.sum(axis=0)   # sum each column → [5, 7, 9]
m.sum(axis=1)   # sum each row    → [6, 15]
m.max(axis=1)   # max per row     → [3, 6]

ufuncs.py

5 — Broadcasting

Broadcasting is NumPy's mechanism for applying operations between arrays of different shapes without making explicit copies of the data. Two shapes are compatible if, when right-aligned, each pair of dimensions is either equal or one of them is 1.

import numpy as np

# Rule: dimensions are compatible if they are equal OR one of them is 1
# NumPy pads dimensions on the LEFT with size-1 until shapes match

a = np.array([[1], [2], [3]])       # shape (3, 1)
b = np.array([10, 20, 30, 40])      # shape (4,) → padded to (1, 4)

a + b
# shape (3, 4):
# [[11, 21, 31, 41],
#  [12, 22, 32, 42],
#  [13, 23, 33, 43]]

# ── Practical examples ──
rng = np.random.default_rng(seed=42)

# Subtract the column mean from each column (zero-centering)
data = rng.random((100, 4))
data_centered = data - data.mean(axis=0)   # (100,4) - (4,) → broadcast

# Add a bias term to each row
weights = np.array([0.1, 0.2, 0.3, 0.4])  # shape (4,)
bias    = 1.0                               # scalar
output  = data @ weights + bias            # (100,4)@(4,) = (100,) + scalar

# Scale each column by a different factor
scale = np.array([2, 0.5, 10, 1])          # shape (4,)
data_scaled = data * scale                  # (100,4) * (4,) → broadcast

broadcasting.py

Shape alignment rule — right-align dimensions and pad missing ones on the left with 1:

Shape alignment (right-align, pad left with 1s):

  a  →  (3, 1)
  b  →  (   4)  →  becomes  (1, 4)
        ------
  out → (3, 4)  ← output shape

Each output[i, j] = a[i, 0] + b[0, j]
No data is copied — NumPy "pretends" the size-1 dimension repeats.

⚠️ Silent Shape Mistakes
Broadcasting succeeds silently even when a shape mismatch is an accident. Always verify unexpected results with print(result.shape) during development.

6 — Linear Algebra

NumPy's np.linalg sub-module wraps LAPACK and BLAS routines, giving you production-grade linear algebra in a single import.

import numpy as np

A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])

# ── Matrix operations ──
A @ B                     # matrix multiply (preferred over np.dot for 2-D+)
np.dot(A, B)              # equivalent
A.T                       # transpose
np.linalg.inv(A)          # inverse  [[-2, 1], [1.5, -0.5]]
np.linalg.det(A)          # determinant → -2.0
np.linalg.matrix_rank(A)  # 2

# ── Solving linear systems: Ax = b ──
b = np.array([1, 2])
x = np.linalg.solve(A, b)   # faster & more stable than inv(A) @ b
# verify: A @ x ≈ b

# ── Decompositions ──
eigenvalues, eigenvectors = np.linalg.eig(A)
U, S, Vt = np.linalg.svd(A)          # singular value decomposition
Q, R     = np.linalg.qr(A)           # QR decomposition
L        = np.linalg.cholesky(A.T @ A)  # Cholesky (requires positive-definite)

# ── Norms ──
np.linalg.norm(A)               # Frobenius norm
np.linalg.norm(A, ord=1)        # max column sum
np.linalg.norm(A, ord=np.inf)   # max row sum

# ── Dot products (1-D vectors) ──
u = np.array([1.0, 2.0, 3.0])
v = np.array([4.0, 5.0, 6.0])
np.dot(u, v)           # 32.0        (u·v scalar)
np.cross(u, v)         # [-3, 6, -3] (u×v vector)
np.outer(u, v)         # outer product (3×3 matrix)

linear_algebra.py

Operation	Function	Notes
Matrix multiply	`A @ B`	Use `@` (PEP 465) for readability
Inverse	`np.linalg.inv(A)`	Avoid for solving systems — use `solve` instead
Solve Ax = b	`np.linalg.solve(A, b)`	More numerically stable than computing the inverse
Eigenvalues	`np.linalg.eig(A)`	Returns (values, vectors); use `eigh` for symmetric matrices
SVD	`np.linalg.svd(A)`	Foundation of PCA, LSA, and recommender systems
Norm	`np.linalg.norm(A)`	Default is Frobenius norm for matrices, L2 for vectors

7 — Practical Patterns & Performance

import numpy as np

# ── Avoid Python loops: use vectorised ops ──

# ❌ slow
def slow_norm(arr):
    total = 0
    for x in arr:
        total += x * x
    return total ** 0.5

# ✅ fast
def fast_norm(arr):
    return np.sqrt(np.dot(arr, arr))   # or np.linalg.norm(arr)


# ── np.vectorize (convenience, NOT speed) ──
def my_func(x):
    return x ** 2 + 1 if x > 0 else 0

vfunc = np.vectorize(my_func)     # still a Python loop under the hood
# Use only for prototyping — replace with where/clip/ufuncs for speed


# ── np.where for conditional vectorisation ──
arr    = np.linspace(-3, 3, 7)
result = np.where(arr > 0, arr ** 2 + 1, 0)   # fully vectorised


# ── Avoid repeated memory allocation in loops ──
# ❌ slow: creates a brand-new array every iteration
out = np.zeros(100)
for i in range(100):
    out = out + np.ones(100) * i

# ✅ use in-place operations
out  = np.zeros(100)
temp = np.ones(100)
for i in range(100):
    out += temp * i     # or: np.add(out, temp * i, out=out)


# ── Profiling with %timeit (Jupyter / IPython) ──
# %timeit slow_norm(arr)
# %timeit fast_norm(arr)


# ── Saving & loading ──
rng = np.random.default_rng(42)
arr = rng.random((1000, 1000))

np.save("data.npy", arr)                                        # binary .npy
arr2 = np.load("data.npy")

np.savetxt("data.csv", arr[:5, :5], delimiter=",", fmt="%.4f") # text CSV
arr3 = np.loadtxt("data.csv", delimiter=",")

# Multiple arrays in one archive
np.savez("arrays.npz", a=arr, b=arr[:10])
loaded = np.load("arrays.npz")
print(loaded["a"].shape)   # (1000, 1000)

performance_patterns.py

💡 Large Arrays That Don't Fit in Memory
For datasets too large to load at once, use np.memmap (memory-mapped file) to treat a binary file on disk as a NumPy array — only the slices you access are loaded into RAM. For structured / multi-dataset storage, prefer h5py (HDF5 format), which supports compression, chunking, and partial reads natively.

Structured Arrays

Structured arrays let you store heterogeneous data (like a CSV with mixed types) in a single NumPy array — each element is a record with named fields.

import numpy as np

# Define a dtype with named fields
dt = np.dtype([
    ("name",  "U20"),     # Unicode string, max 20 chars
    ("age",   np.int32),
    ("score", np.float64),
])

data = np.array([
    ("Alice", 30, 95.5),
    ("Bob",   25, 87.2),
    ("Carol", 28, 91.0),
], dtype=dt)

# Access by field name
print(data["name"])    # ['Alice' 'Bob' 'Carol']
print(data["score"])   # [95.5 87.2 91. ]

# Filter with boolean indexing
data[data["score"] > 90]   # Alice and Carol

# Sort by field
data_sorted = np.sort(data, order="score")   # ascending by score

structured_arrays.py

For most tabular data, Pandas DataFrames (Lesson 44) are more ergonomic than structured arrays — they offer named columns, richer indexing, and many built-in operations. Use structured arrays when you need raw NumPy speed on mixed-type records or when integrating with C extensions.

Stacking, Splitting & Combining

import numpy as np

a = np.array([[1, 2], [3, 4]])   # (2, 2)
b = np.array([[5, 6], [7, 8]])   # (2, 2)

# ── Stacking ──
np.vstack([a, b])     # vertical stack  → (4, 2)
np.hstack([a, b])     # horizontal stack → (2, 4)
np.dstack([a, b])     # depth stack      → (2, 2, 2)
np.concatenate([a, b], axis=0)  # same as vstack
np.concatenate([a, b], axis=1)  # same as hstack
np.stack([a, b], axis=0)        # new axis → (2, 2, 2)
np.stack([a, b], axis=2)        # new axis at position 2 → (2, 2, 2)

# ── Splitting ──
c = np.arange(12).reshape(4, 3)
np.vsplit(c, 2)         # split into 2 equal parts along axis 0
np.hsplit(c, 3)         # split into 3 equal parts along axis 1
np.array_split(c, 3, axis=0)   # unequal splits allowed

# ── Repeat & tile ──
np.repeat(np.array([1, 2, 3]), 3)       # [1,1,1,2,2,2,3,3,3]
np.tile(np.array([1, 2, 3]), (2, 3))    # 2 rows, repeat 3 times along cols

# ── Sorting & searching ──
arr = np.array([3, 1, 4, 1, 5, 9, 2, 6])
np.sort(arr)           # sorted copy
arr.sort()             # in-place sort
np.argsort(arr)        # indices that would sort the array
np.searchsorted(np.sort(arr), 5)  # binary search → insertion index
np.unique(arr)         # unique values (sorted)
np.unique(arr, return_counts=True)  # with occurrence counts

stacking.py

A Complete NumPy Workflow

Putting it all together: load tabular data, clean it, compute statistics, and apply a vectorised transformation.

import numpy as np

# ── Simulate loading a dataset ──
rng = np.random.default_rng(42)
raw = rng.normal(loc=50, scale=15, size=(200, 4))   # 200 samples, 4 features

# ── Inspect ──
print(raw.shape)     # (200, 4)
print(raw.dtype)     # float64
print(raw[:3])       # first 3 rows

# ── Descriptive statistics per feature ──
print("mean  :", raw.mean(axis=0))
print("std   :", raw.std(axis=0))
print("min   :", raw.min(axis=0))
print("max   :", raw.max(axis=0))

# ── Min-max normalisation (vectorised, no loops) ──
col_min = raw.min(axis=0)   # shape (4,)
col_max = raw.max(axis=0)   # shape (4,)
normalised = (raw - col_min) / (col_max - col_min)   # broadcasting

# ── Z-score standardisation ──
mean = raw.mean(axis=0)
std  = raw.std(axis=0)
standardised = (raw - mean) / std

# ── Remove rows where any feature > 3 std devs from mean ──
outlier_mask = np.any(np.abs(standardised) > 3, axis=1)
clean = raw[~outlier_mask]
print(f"Removed {outlier_mask.sum()} outlier rows; {len(clean)} remain")

# ── Correlation matrix ──
corr = np.corrcoef(clean.T)   # (4, 4) correlation matrix
print("Feature correlations:\n", np.round(corr, 2))

# ── Save clean dataset ──
np.save("clean_data.npy", clean)

workflow.py

Best Practices

Never loop over NumPy arrays in Python — always look for a vectorised operation, ufunc, or np.where / np.select equivalent first.
Choose the right dtype — float32 uses half the memory of float64 and is faster on GPUs; use int8/uint8 for image pixel data.
Slices are views, not copies — call .copy() explicitly when you need an independent array; modifying a view modifies the original.
Verify shapes with assertions during development — assert result.shape == (100, 4), result.shape. Silent broadcasting errors are hard to debug.
Use np.linalg.solve(A, b) not np.linalg.inv(A) @ b — solve is faster and numerically more stable.
Seed the RNG for reproducibility — use np.random.default_rng(seed=42) (not the legacy np.random.seed()) in any experiment.
Use np.savez_compressed for large arrays — it compresses automatically and can store multiple arrays in a single file.
Profile before optimising — use %timeit in Jupyter or cProfile to find the actual bottleneck before rewriting.

Exercises

Exercise 1 — Vectorised Statistics from Scratch

Implement the following without any Python loops, using only NumPy:

row_normalise(X) — divide each row by its L2 norm so every row has unit length.
pairwise_distances(X) — return the N×N matrix of Euclidean distances between all row pairs of an (N, D) array. Use broadcasting; no scipy.spatial.
softmax(x) — compute exp(x) / sum(exp(x)) in a numerically stable way (subtract the max before exponentiation).
Test all three on small inputs and verify with known values.

💡 Hint

def row_normalise(X: np.ndarray) -> np.ndarray:
    norms = np.linalg.norm(X, axis=1, keepdims=True)   # (N, 1)
    return X / norms

def pairwise_distances(X: np.ndarray) -> np.ndarray:
    # ||a - b||² = ||a||² + ||b||² - 2 a·b
    sq = np.sum(X ** 2, axis=1)                    # (N,)
    dist_sq = sq[:, None] + sq[None, :] - 2 * (X @ X.T)
    return np.sqrt(np.clip(dist_sq, 0, None))      # clip negatives from floating point

def softmax(x: np.ndarray) -> np.ndarray:
    e = np.exp(x - x.max())   # subtract max for numerical stability
    return e / e.sum()

Exercise 2 — Image Processing with NumPy

Images are just NumPy arrays (H × W × 3 for RGB). Implement these operations without any image library — pure NumPy only:

Create a synthetic 100×100 RGB image where the red channel is a horizontal gradient (0 → 255), green is a vertical gradient, and blue is constant 128.
Convert to greyscale using the luminosity formula: 0.2126·R + 0.7152·G + 0.0722·B.
Apply a horizontal flip and a vertical flip.
Crop the centre 60×60 pixels.
Normalise pixel values to [0, 1] as float32.

💡 Hint

import numpy as np

H, W = 100, 100
img = np.zeros((H, W, 3), dtype=np.uint8)
img[:, :, 0] = np.linspace(0, 255, W, dtype=np.uint8)            # R: horizontal
img[:, :, 1] = np.linspace(0, 255, H, dtype=np.uint8)[:, None]   # G: vertical
img[:, :, 2] = 128                                                 # B: constant

# Greyscale
weights = np.array([0.2126, 0.7152, 0.0722])
grey = (img @ weights).astype(np.uint8)   # (H, W)

# Flips
h_flip = img[:, ::-1, :]
v_flip = img[::-1, :, :]

# Crop centre 60×60
cy, cx = H // 2, W // 2
crop = img[cy-30:cy+30, cx-30:cx+30, :]

# Normalise
norm = img.astype(np.float32) / 255.0

Exercise 3 — Linear Regression with NumPy

Implement ordinary least-squares linear regression from scratch using only NumPy:

Generate synthetic data: y = 3x + 2 + noise with 100 points.
Implement fit(X, y) that computes the closed-form OLS solution: w = (XᵀX)⁻¹Xᵀy using np.linalg.solve.
Implement predict(X, w) and mse(y_true, y_pred).
Implement fit_gradient_descent(X, y, lr=0.01, epochs=1000) and verify it converges to the same weights as the closed-form solution.
Plot the data and the fitted line (optional if matplotlib not installed).

💡 Hint

import numpy as np

rng = np.random.default_rng(42)
X_raw = rng.uniform(0, 10, 100)
y = 3 * X_raw + 2 + rng.normal(0, 1, 100)

# Add bias column
X = np.column_stack([np.ones(100), X_raw])   # (100, 2)

def fit(X, y):
    # w = (XᵀX)⁻¹Xᵀy — use solve for numerical stability
    return np.linalg.solve(X.T @ X, X.T @ y)

def predict(X, w):
    return X @ w

def mse(y_true, y_pred):
    return np.mean((y_true - y_pred) ** 2)

w = fit(X, y)
print(f"Intercept: {w[0]:.3f}, Slope: {w[1]:.3f}")  # ≈ 2, 3

def fit_gd(X, y, lr=0.01, epochs=1000):
    w = np.zeros(X.shape[1])
    n = len(y)
    for _ in range(epochs):
        grad = X.T @ (X @ w - y) / n
        w -= lr * grad
    return w

w_gd = fit_gd(X, y)
print(f"GD  — Intercept: {w_gd[0]:.3f}, Slope: {w_gd[1]:.3f}")