AskAiml
Sign in Get started

NumPy essentials — vectorise everything

Array shapes, broadcasting, views vs copies, and the dtypes that decide your memory footprint.

📊 Module 5 8 min read Not started

Why this matters

NumPy is the foundation of every Python data / ML / scientific library — pandas, scikit-learn, PyTorch, JAX, OpenCV all build on ndarray. Vectorised NumPy is 10–1000× faster than pure Python loops because the work happens in tight C code, releasing the GIL. This lesson teaches you to think in arrays.

Learning objectives

  1. Create and reshape ndarrays.
  2. Slice, index, and boolean-mask arrays.
  3. Apply broadcasting correctly.
  4. Vectorise instead of looping.
  5. Use numpy.random, linalg, and reductions.

1. The ndarray

python
import numpy as np

a = np.array([1, 2, 3, 4])
a.shape          # (4,)
a.dtype          # dtype('int64')
a.ndim           # 1
a.size           # 4
a.itemsize       # 8 (bytes per element)
a.nbytes         # 32

NumPy arrays are:

  • Homogeneous (one dtype).
  • Contiguous in memory.
  • Fixed-shape at creation (resizing is rare).

Common constructors

python
np.array([[1,2],[3,4]])
np.zeros((3, 4))
np.ones((2, 3), dtype=np.float32)
np.full((3,), 7.0)
np.empty((2, 2))                # uninitialised; faster, may contain garbage
np.eye(3)                       # identity
np.arange(0, 10, 2)             # [0, 2, 4, 6, 8]
np.linspace(0, 1, 5)            # 5 evenly-spaced points incl. endpoints
np.logspace(0, 3, 4)            # [1, 10, 100, 1000]

# From random
rng = np.random.default_rng(seed=42)
rng.standard_normal((3, 3))
rng.integers(0, 100, size=10)
rng.uniform(0, 1, size=5)
rng.choice([1,2,3], size=10, p=[0.5, 0.3, 0.2])

Always seed the Generator (default_rng), not the legacy np.random.* functions.

dtypes

dtypeuse
int8, int16, int32, int64signed integers
uint8, ..., uint64unsigned (uint8 for images)
float16, float32, float64floats (float32 for ML)
bool_True/False
complex64, complex128complex
str_ (fixed-width) / objectstrings (slow; avoid)
python
a.astype(np.float32)

2. Indexing and slicing

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

a[0, 1]                       # 1
a[1]                          # [4,5,6,7]
a[:, 0]                       # [0,4,8] (first column)
a[1:, 2:]                     # subarray
a[::-1]                       # reversed rows
a[::2, ::2]                   # every other row, every other column

Slices are views (no copy). Mutations propagate to the original.

python
b = a[1]
b[0] = 99
a[1]                          # [99, 5, 6, 7]   — changed!
b = a[1].copy()               # explicit copy if needed

Fancy indexing (always copies)

python
a[[0, 2], [1, 3]]            # diagonal of corners: a[0,1] and a[2,3]
a[[True, False, True]]       # boolean mask; rows where True
a[a > 5]                      # 1D array of values > 5
a[(a > 3) & (a < 9)]          # combine masks with &, |, ~ (parentheses!)

np.where:

python
np.where(a > 5, a, 0)         # element-wise: a if >5 else 0
np.where(a > 5)               # indices where condition true

3. Shape manipulation

python
a = np.arange(12)
a.reshape(3, 4)
a.reshape(3, -1)              # -1 means "infer"
a.reshape(-1, 1)              # column vector

np.concatenate([a, a])        # along axis 0
np.stack([a, a])              # add new axis
np.vstack([a, a]); np.hstack([a, a])
np.split(a, 4)                # 4 equal parts
np.expand_dims(a, axis=0)     # (12,) -> (1, 12)
a[..., np.newaxis]            # add trailing axis
a.T                            # transpose
np.transpose(a, axes=(1, 0))
np.swapaxes(arr, 0, 1)

Contiguity & memory layout

python
a.flags["C_CONTIGUOUS"]       # True if row-major
a.flags["F_CONTIGUOUS"]       # True if column-major
np.ascontiguousarray(a)

For interop with C libraries (PyTorch, OpenCV) and best cache behaviour, use C-contiguous arrays.

4. Broadcasting — the killer feature

NumPy automatically aligns shapes for element-wise ops:

python
a = np.array([[1,2,3], [4,5,6]])      # shape (2, 3)
b = np.array([10, 20, 30])             # shape (3,)
a + b                                  # broadcasts b across rows: [[11,22,33],[14,25,36]]

c = np.array([[100], [200]])           # shape (2, 1)
a + c                                  # broadcasts across columns

Rules

  1. Align shapes from the trailing dimension.
  2. Two dims are compatible if equal or one is 1.
  3. Otherwise → ValueError.
a: (2, 3)
b: (   3)
   → align as (1, 3) and (2, 3) → result (2, 3)

a: (4, 3)
b: (4, 1)
   → result (4, 3)

Broadcasting avoids explicit loops AND avoids creating intermediate copies. Reading shapes mentally is a core skill.

5. Vectorisation

Pure Python:

python
result = [x*x + 1 for x in data]

NumPy:

python
result = data ** 2 + 1

100× faster for size ≥ 1000.

Common reductions

python
a.sum(); a.mean(); a.std(); a.min(); a.max()
a.sum(axis=0)         # column sums
a.sum(axis=1)         # row sums
a.sum(axis=0, keepdims=True)   # keeps shape (1, n)

np.cumsum(a); np.cumprod(a)
np.percentile(a, 95)
np.median(a)
np.var(a, ddof=1)     # sample variance
np.argmin(a); np.argmax(a)
a.any(); a.all()

Element-wise functions (ufuncs)

python
np.exp(a), np.log(a), np.sqrt(a), np.sin(a)
np.maximum(a, b)   # element-wise max
np.minimum(a, b)
np.clip(a, 0, 10)
np.abs(a)
np.where(cond, x, y)
np.isnan(a); np.isfinite(a)

All operate element-wise, broadcast naturally, and run in C.

6. Linear algebra

python
np.dot(a, b)            # matrix multiplication (use @)
a @ b                   # preferred
np.linalg.inv(a)
np.linalg.det(a)
np.linalg.eig(a)         # eigenvalues, eigenvectors
np.linalg.svd(a)         # SVD
np.linalg.solve(A, b)    # solve Ax = b
np.linalg.norm(a)        # Frobenius / L2
np.linalg.matrix_rank(a)
np.linalg.pinv(a)        # pseudoinverse

For ML, prefer np.linalg.solve(A, b) over inv(A) @ b — numerically more stable.

7. Boolean and integer masking idioms

python
data = rng.standard_normal(1000)
positive = data[data > 0]
data[data < 0] = 0
data = np.clip(data, 0, None)

# Count true
(data > 0).sum()
np.count_nonzero(data > 0)

# Pick top k indices
top_k = np.argpartition(data, -10)[-10:]      # not sorted
top_k_sorted = top_k[np.argsort(-data[top_k])]

# argsort tie-broken indices
np.argsort(data)

8. Saving and loading

python
np.save("arr.npy", a)
b = np.load("arr.npy")

np.savez("arrs.npz", a=a, b=b)         # multiple arrays
data = np.load("arrs.npz"); data["a"]

# Compressed
np.savez_compressed("arrs.npz", a=a)

For very large arrays use memory-mapped files:

python
m = np.memmap("big.dat", dtype="float32", mode="w+", shape=(10_000_000,))
m[:] = ...                                # writes to disk
del m                                       # flushes

For interop with pandas / parquet / etc., go through pyarrow.

9. Random numbers (the modern API)

python
rng = np.random.default_rng(42)
rng.normal(loc=0, scale=1, size=10)
rng.uniform(0, 1, size=10)
rng.integers(0, 100, size=10)
rng.choice([1, 2, 3], size=5, replace=True, p=[0.6, 0.3, 0.1])
rng.shuffle(a)            # in-place
rng.permutation(a)        # copy

Pass a Generator to functions that need randomness — keeps tests deterministic.

10. Performance tips

  • Avoid Python loops over arrays. If you find yourself writing for i in range(len(a)), find a vectorised expression.
  • Use views (a[:, 0]) over copies when possible.
  • Watch dtypes: float64 + float32 upcasts to float64; uint8 - 1 wraps to 255. Explicit .astype when intent matters.
  • Pre-allocate: build outputs with np.empty and fill in, rather than np.append in a loop.
  • np.einsum for elegant tensor contractions:
    python
    np.einsum("ij,jk->ik", A, B)              # matrix multiply
np.einsum("ii->i", A)                     # diagonal
np.einsum("bij,bjk->bik", A, B)           # batched matmul
  • numexpr (pip install numexpr) can sometimes outperform NumPy on big expressions by avoiding intermediates.

11. Worked example: image manipulation

python
import numpy as np
from PIL import Image

img = np.asarray(Image.open("photo.jpg"))         # shape (H, W, 3), dtype uint8

# Convert to grayscale (Rec. 709 luminance)
gray = (img @ np.array([0.2126, 0.7152, 0.0722])).astype(np.uint8)

# Invert
inverted = 255 - img

# Crop centre 200x200
h, w = img.shape[:2]
crop = img[h//2 - 100 : h//2 + 100, w//2 - 100 : w//2 + 100]

# Brightness boost (clip to avoid overflow)
brighter = np.clip(img.astype(np.int32) + 30, 0, 255).astype(np.uint8)

Image.fromarray(gray).save("gray.png")

No loops. Each operation runs at C speed.

Hands-on lab (2 hours)

  1. Generate 1M normally distributed numbers; compute mean, std, 95th percentile.
  2. Reshape np.arange(24) into (2, 3, 4); sum along each axis; print shapes.
  3. Build a 5×5 multiplication table with broadcasting (np.arange(1,6)[:, None] * np.arange(1,6)).
  4. Implement min-max scaling on a column matrix without loops.
  5. Solve Ax = b with np.linalg.solve for a 1000×1000 system.
  6. Convert a colour image to greyscale (above example); time vs Python-loop version.
  7. Bonus: implement k-means initialisation (k-means++) using np.argmin + cumulative probability sampling.

Common pitfalls

  1. Confusing view vs copy; mutating a view changes the original.
  2. Forgetting axis= — a.sum() returns scalar, a.sum(axis=0) returns array.
  3. == between arrays returns an array; use .all() / .any() for booleans.
  4. Mixing np.array (multi-dim) with np.matrix (deprecated; never use).
  5. Using lists where arrays would work — silently slow.
  6. np.append / np.concatenate in a loop — O(n²). Pre-allocate.
  7. Reshaping non-contiguous arrays with implicit copy can spike memory.

Self-check

  1. Why is NumPy fast?
  2. State the broadcasting rules.
  3. Difference between a view and a copy.
  4. Why use default_rng over np.random.*?
  5. When use solve vs inv?

References

  • Python for Data Analysis, Wes McKinney — Chapter 4.
  • From Python to NumPy, Nicolas Rougier (free book).
  • NumPy user guide: https://numpy.org/doc/stable/user/.
  • Travis Oliphant, Guide to NumPy.
  • "100 NumPy exercises" — github.com/rougier/numpy-100.

Sign in to save your progress and earn badges.