Is it ever fine to just write a normal Python loop when working with data?

Yes — a loop over a small collection, or one whose body does I/O (an API call, a file write) rather than numeric computation, is rarely the bottleneck and does not need to be rewritten. The concern is specifically large arrays with simple numeric work in the loop body, which is where Python's per-iteration overhead actually dominates.

Does `numpy.vectorize` give the same speedup as writing a real vectorized expression?

No — despite the name, `numpy.vectorize` is a convenience wrapper that still calls a Python function once per element internally, so it does not remove the interpreter overhead a true vectorized expression (using array arithmetic and NumPy's built-in functions directly) avoids.

Does using Jupyter notebooks at all count as bad practice?

No — notebooks are the right tool specifically for exploration, where re-running cells out of order and iterating quickly is exactly what you want. The practice worth avoiding is treating a notebook as the final, trustworthy record of a result meant to be reproduced or shared — that is what scripts and pinned environments are for.

Do I need to know both pandas and NumPy, or is one enough?

Both, in practice — pandas is built on top of NumPy and is the more natural tool for labeled, tabular data (loading, cleaning, joining), while NumPy is what you reach for once data is already numeric arrays, especially for anything feeding directly into a model. Most real preprocessing pipelines use both together.

beginnerintermediateadvanced

AI Full-Stack · Section 3

Python for AI

Level: beginner
Read: 45 min
Concepts: 5

The Python subset AI work reaches for constantly: the language features ML libraries are built around (generators, context managers, decorators, typing), NumPy arrays and vectorization — the performance model every numerical library assumes — Pandas for getting tabular data feature-ready, notebooks and environments for exploration versus reproducibility, and the judgment to know when a plain Python loop has become the bottleneck.

What is true here

ML code leans hardest on generators, context managers, decorators and type hints
Vectorization pushes a loop into compiled code — 50-100x faster than the Python equivalent on large arrays
Broadcasting lets NumPy operate on differently-shaped arrays without explicit loops or copies
Pandas turns raw tabular files into feature-ready data: load, clean, join, groupby, handle missing values
A notebook is for exploration, not a reliable record — reproducibility needs code, pinned deps, data version and seed together

What you will be able to do

Recognize which Python features (generators, context managers, decorators) a piece of ML code is using and why
Predict whether two array shapes will broadcast together, and why
Rewrite a Python loop over array data as a vectorized NumPy/pandas expression
Join, group, and handle missing values in a pandas DataFrame
Explain the four ingredients a reproducible experiment needs

Python, NumPy and Pandas for ML work

The language and library subset that shows up in nearly every ML codebase, from data loading through the performance model every numerical library assumes.

Python fundamentals AI work leans on hardest

standardbeginner

AI code is Python code, and a handful of Python features show up constantly in ML libraries: generators for streaming large datasets without loading them all into memory, decorators for wrapping training steps, context managers for cleanly acquiring and releasing resources like a database connection or a GPU context, and typing for catching shape and schema mistakes before runtime.

Think of it as

None of this is AI-specific — it is the same Python every backend engineer uses — but ML code reaches for a specific subset constantly. A dataset that does not fit in memory is streamed with a generator (`yield` one example at a time) rather than built as one giant list. A context manager (`with ...:`) shows up whenever a resource needs guaranteed cleanup — a file, a database session, a `torch.no_grad()` block that turns off gradient tracking for inference. Decorators wrap a function with extra behavior without changing its body — `@torch.no_grad()`, `@retry`, `@app.route` in a serving API. Typing (type hints) catches an entire class of bugs before they reach production: passing a list where a NumPy array was expected, or an int where a tensor's shape expects a tuple.

python

def batches(data, size):
    for i in range(0, len(data), size):
        yield data[i:i + size]          # generator: one batch at a time

with torch.no_grad():                    # context manager: no gradient tracking
    predictions = model(inputs)

@retry(times=3)                          # decorator: wraps a function
def call_model_api(prompt: str) -> str:  # type hints on the signature
    ...

Remember: ML code leans hardest on generators (stream data), context managers (guaranteed cleanup), decorators (wrap behavior) and type hints (catch shape/type bugs early) — the same Python, applied to the same recurring shapes of problem.

NumPy arrays, broadcasting and vectorization

corebeginner

A NumPy array is a grid of numbers, all the same type, stored contiguously in memory so operations on it run at C speed instead of Python speed. Vectorization means expressing an operation over the whole array at once instead of writing a Python `for` loop over its elements.

Think of it as

A Python `for` loop over a list pays Python's per-element overhead on every iteration. A vectorized NumPy operation pushes the whole loop down into compiled C code, running orders of magnitude faster on the same data — this is why `array_a + array_b` is fast but `[a[i] + b[i] for i in range(len(a))]` on the same arrays is slow. Broadcasting is the rule that lets NumPy operate on arrays of different shapes without you writing explicit loops or copies: a smaller array is conceptually 'stretched' to match a larger one's shape, following strict compatibility rules, and the actual computation never materializes the stretched copy. Every array has a dtype (float32, int64, bool...) that fixes both its precision and its memory footprint — float32 uses half the memory of float64, which matters a great deal once arrays reach GPU scale. Reshaping changes how the same underlying data is viewed (a flat 12-element array becomes a 3×4 matrix) without copying it, as long as the memory stays contiguous; a transpose or certain slices break that contiguity and force a real copy the next time an operation needs one.

python

import numpy as np

a = np.array([1, 2, 3], dtype=np.float32)   # dtype fixes precision + memory
b = a.reshape(3, 1)                          # view, not a copy (contiguous)
c = a + np.array([[10], [20]])               # broadcasting: (3,) + (2,1) -> (2,3)
mask = a > 1                                 # boolean mask
filtered = a[mask]                           # fancy indexing -> always a copy

What we're doing: Compare a Python loop to a vectorized NumPy operation computing the same result.

vectorize_vs_loop.pypython

import numpy as np
import time

n = 1_000_000
a = list(range(n))
b = list(range(n))

start = time.perf_counter()
result_loop = [a[i] + b[i] for i in range(n)]
loop_time = time.perf_counter() - start

arr_a = np.arange(n)
arr_b = np.arange(n)

start = time.perf_counter()
result_vec = arr_a + arr_b
vec_time = time.perf_counter() - start

print(f"loop:       {loop_time:.4f}s")
print(f"vectorized: {vec_time:.4f}s  ({loop_time / vec_time:.0f}x faster)")

8: A pure-Python list comprehension pays interpreter overhead on every one of a million iterations — this is the baseline every ML library is built to avoid.
17: The vectorized version expresses the same computation as one operation over the whole array, executed in compiled code — the speedup on a million elements is typically 50-100x.
20: This gap is exactly why every ML library — NumPy, pandas, PyTorch, TensorFlow — pushes you toward array operations and treats an explicit Python loop over elements as something to avoid.

Why this works: The performance gap between a Python loop and a vectorized operation is not a minor tuning detail — it is the reason every serious numerical library in Python is built the way it is, and not understanding it leads directly to code that is 50-100x slower than it needs to be.

Looping over a NumPy array element by element

Wrong

python

normalized = np.zeros_like(arr)
for i in range(len(arr)):
    normalized[i] = (arr[i] - mean) / std
# Works, but re-introduces Python's per-element
# overhead on top of a library built to avoid it.

Better

python

normalized = (arr - mean) / std
# One vectorized expression. Broadcasting handles
# "mean" and "std" being scalars applied to every
# element — no loop, no per-element Python overhead.

What you see: Preprocessing that should take milliseconds on a few thousand rows instead takes seconds, and profiling shows nearly all the time inside a Python `for` loop over array indices rather than inside NumPy itself.

Why: Indexing a NumPy array inside a Python loop (`arr[i]`) pays Python's interpreter overhead on every access, throwing away the entire reason NumPy arrays exist. Any operation expressible as array arithmetic, broadcasting, or a NumPy function should be — the loop is almost always a sign the vectorized equivalent was not found.

Broadcasting compatibility, by shape

Broadcasting compatibility, by shape
Shape A	Shape B	Result	Why
(3, 4)	(4,)	(3, 4)	B is stretched across every row of A
(3, 1)	(1, 4)	(3, 4)	Both stretch to fill the missing dimension
(3, 4)	(3,)	Error	Trailing dimensions (4 vs 3) do not match and neither is 1
(1000, 512)	(512,)	(1000, 512)	A per-feature vector applied to every one of 1000 rows

Remember: Vectorization pushes a loop into compiled code — 50-100x faster than the Python equivalent. Broadcasting stretches smaller-shaped arrays to match larger ones by fixed rules. Basic slicing is a free view; fancy indexing and some reshapes copy. dtype fixes both precision and memory.

Pandas for data loading and cleaning

standardbeginner

Pandas is a library for working with tabular data — rows and named columns, like a spreadsheet you can script. It handles loading data from files, cleaning it, joining separate tables together, grouping rows to compute aggregates, and finding and handling missing values, all before that data becomes model features.

Think of it as

Think of a pandas DataFrame as the step between raw files and a NumPy array a model can train on. Loading brings a CSV, database table, or JSON file into a DataFrame. Cleaning fixes inconsistent types, duplicate rows, and malformed values. A join combines two tables on a shared key, the same operation as a SQL join — critical when features live in separate tables (user profile in one file, transaction history in another). Grouping (`groupby`) computes an aggregate per category, like average purchase per customer. Missing values are pervasive in real data, and pandas gives you the tools to find them (`isna()`), and either fill them with a sensible value or drop the rows — a decision that belongs to feature engineering, not to pandas itself. The whole library exists to get data into a clean, numeric, feature-ready shape before it reaches a model.

python

import pandas as pd

df = pd.read_csv("transactions.csv")
merged = df.merge(users, on="user_id", how="left")   # join on a key
per_user = merged.groupby("user_id")["amount"].sum()  # aggregate per group

missing = df.isna().sum()                # count missing per column
df["amount"] = df["amount"].fillna(0)    # fill, a modeling decision
df = df.dropna(subset=["user_id"])       # or drop rows with no key at all

Remember: Pandas gets tabular data from raw files to feature-ready: load, clean, join tables on a shared key, groupby for per-category aggregates, and handle missing values deliberately — fill or drop is a modeling decision, not a default.

Notebooks, environments and reproducible experiments

standardbeginner

A notebook (Jupyter) lets you run Python in small, re-orderable cells and see results inline — good for exploring data, bad for production code. A Python environment isolates one project's package versions from every other project's, so installing one library never silently breaks another.

Think of it as

Notebooks and environments solve different problems, and conflating them causes real damage. A notebook's cells can be run out of order, which is exactly what makes exploration fast — and exactly what makes a notebook's final state untrustworthy as a record of what actually happened, since a cell run three times with edits in between leaves no trace of that history. That is why production and shared pipelines are refactored out of notebooks into ordinary scripts once they stabilize. An environment (a virtualenv, conda env, or similar) pins exactly which package versions a project uses, isolated from every other project on the machine — without it, upgrading a library for one project can silently break another that assumed an older API. Reproducibility for an experiment means someone else — or you, in six months — can rerun it and get the same result: the same code, the same pinned dependencies, the same data version, and the same random seed, all recorded together.

bash

python -m venv .venv && source .venv/bin/activate   # isolated environment
pip install -r requirements.txt                       # pinned dependencies

jupyter notebook                                       # exploration
# once code stabilizes: move it out of the notebook
# into a script/module that runs top to bottom

Remember: A notebook's re-orderable cells are great for exploration and unreliable as a record — move stable code into scripts. Reproducibility needs code version, pinned dependencies, data version and random seed together, not any one alone.

When a Python loop is too slow

standardbeginner

A Python `for` loop is fine for a few thousand iterations of lightweight work. It becomes a real bottleneck once it runs millions of times, or once each iteration does numeric work that a vectorized library could do instead — that is the signal to reach for NumPy, pandas, or a native library rather than a loop.

Think of it as

The decision is about scale and what kind of work is inside the loop, not a blanket rule against loops. Looping over a short list to call an API or write a log line is fine — the loop overhead is nothing next to the work each iteration does. Looping over a million numeric values to do arithmetic is where Python's per-iteration interpreter overhead dominates the actual computation, and that is the signature to profile for: a loop whose body is simple numeric or array work, run a very large number of times. The fix is almost always to reframe the problem as an array operation — a vectorized NumPy/pandas expression, or a call into a native library (compiled C, Rust, or a GPU kernel) that does the same loop internally, without Python's overhead on each step.

text

loop over a SHORT list, body does I/O        -> loop is fine
loop over a LARGE array, body is numeric math  -> vectorize it
  1. profile first (don't guess)
  2. reframe as an array operation (NumPy/pandas)
  3. if still too slow: a native/compiled library

Remember: A loop is a problem when it runs many times AND its body is simple numeric work — profile to confirm, then reframe as a vectorized operation or hand it to a native library instead of micro-optimizing the loop.