Case Study 1: The Notebook That Wouldn't Finish

"The profiler is not there to confirm your guess about where the time goes. It is there to correct it."

Executive Summary

A researcher has a Python notebook that simulates heat spreading along a one-dimensional rod. It is readable, it is correct, and it is unusably slow: a run that should take a second takes most of a minute, and the grid they actually need is ten times larger. This case study walks the exact diagnosis-and-repair a computational scientist performs dozens of times a year. We locate the hot loop, establish why it cannot simply be vectorized away in NumPy, port that one loop to Fortran, wrap it with f2py, drive the simulation from the original Python, and validate the result against a known steady state. The notebook stays a notebook; only its engine changes. By the end you will be able to take any slow numerical Python script and decide — with evidence, not vibes — which ten lines to move to Fortran and how.

Skills applied: the two-language workflow and locating the hot loop (§15.1); wrapping a kernel with f2py (§15.2); dtype and the 1-D-is-both-contiguous shortcut (§15.3); honest benchmarking and same-answer validation (§15.5).

Background

The rod is modeled as a line of n cells. The two ends are held at fixed temperatures (left hot, right cold), and each interior cell relaxes toward its neighbors according to the explicit update

$$ u_i^{\,\text{new}} = u_i + r\,(u_{i-1} - 2u_i + u_{i+1}), $$

repeated for many timesteps. The researcher's Python is a faithful transcription of that formula:

def simulate_py(u, r, nsteps):
    n = len(u)
    for _ in range(nsteps):              # time loop
        u_old = u.copy()
        for i in range(1, n - 1):        # space loop — the hot one
            u[i] = u_old[i] + r * (u_old[i-1] - 2*u_old[i] + u_old[i+1])
    return u

It is a textbook double loop: an outer loop over time and an inner loop over space, with the ends left untouched so they stay fixed. On a rod of a few hundred cells for a few thousand steps, it crawls.

Phase 1 — Locate the Hot Loop

The first discipline of §15.1 is to not guess. Before porting anything, ask where the time actually goes. A quick count is enough to reason about it: the body of the inner loop runs nsteps × (n-2) times. For nsteps = 5000 and n = 500 that is about 2.5 million executions of an interpreted line — each one dispatching Python bytecode, checking types, and boxing floats around a single useful arithmetic expression. Everything else in the notebook (building u, plotting the final profile, printing a summary) runs a handful of times. The 2.5-million-times line is the program.

The reasoning that matters: this is the two-language workflow's central claim (§15.1) applied concretely. Ninety-nine percent of the runtime lives in one inner loop; the other ninety-nine percent of the code barely matters. We do not rewrite the notebook. We move exactly one loop.

Phase 2 — Decide What to Port (and Why Not Just Vectorize)

A NumPy-fluent reader will object: the space loop is a stencil — surely u[1:-1] = u_old[1:-1] + r * (u_old[:-2] - 2*u_old[1:-1] + u_old[2:]) vectorizes it. That is true, and if a single sweep were the whole computation you would stop there. But look again at the structure: the time loop wraps the space loop, and each timestep needs the array the previous timestep produced. That is a loop-carried dependency, and no array expression can flatten it — you must iterate over time no matter what.

So you have two honest choices. Vectorize the inner sweep and keep the Python time loop (fast sweep, but still nsteps Python-level iterations and a full-array temporary each step), or move the whole double loop into Fortran and cross the boundary only once. The second is what a production code does, and it is what we do here: port the entire simulate — both loops — to Fortran, and let Python call it a single time.

Phase 3 — Write and Wrap the Fortran Kernel

The Fortran mirrors the Python exactly, but compiled:

! diffuse.f90 — the whole simulation (both loops) in one wrapped kernel
module diffuse
  use, intrinsic :: iso_fortran_env, only: dp => real64
  implicit none
contains

  subroutine simulate(u, u_out, r, nsteps, n)
    integer,  intent(in)  :: n, nsteps
    real(dp), intent(in)  :: u(n)
    real(dp), intent(out) :: u_out(n)
    real(dp), intent(in)  :: r
!f2py intent(hide), depend(u) :: n = shape(u, 0)
    real(dp) :: old(n)
    integer  :: i, t
    u_out = u                                   ! start from the input; ends stay fixed
    do t = 1, nsteps
      old = u_out
      do i = 2, n - 1
        u_out(i) = old(i) + r * (old(i-1) - 2.0_dp*old(i) + old(i+1))
      end do
    end do
  end subroutine simulate

end module diffuse

$ f2py -c -m difflib diffuse.f90

Because we made simulate return a new array (intent(out) :: u_out) and hid n, the Python signature becomes simulate(u, r, nsteps) -> u_out. Note the whole time loop lives inside Fortran now: Python calls across the boundary exactly once, not nsteps times, so the crossing cost is negligible.

Phase 4 — Get the Data Across Correctly

The field is one-dimensional, which spares us the entire order='F' saga of §15.3: a 1-D array is both C- and F-contiguous, so there is no layout mismatch and no copy. Only the dtype matters. We make sure the array is float64 to match real(dp):

import numpy as np
import difflib

n = 5
u = np.zeros(n, dtype=np.float64)   # float64 matches real(dp); 1-D needs no order='F'
u[0] = 100.0                        # left end hot (held); right end cold (held)

one_step = difflib.diffuse.simulate(u, r=0.25, nsteps=1)
print(one_step)

# Expected output:
# [100.  25.   0.   0.   0.]

Hand-check that one step, because you never trust a ported kernel you have not verified once. With $r = 0.25$ and $u = [100, 0, 0, 0, 0]$, only the ends are fixed; interior cell 2 (Python index 1) updates to $0 + 0.25\,(100 - 2\cdot 0 + 0) = 25$, while cells 3 and 4 see only zeros and stay $0$. So after one step the rod is $[100, 25, 0, 0, 0]$ — exactly what printed. The port is faithful.

Phase 5 — Measure and Validate

Now the two questions that decide whether the port was worth it: is it faster, and is it right?

For speed, time both versions on the real problem size and — non-negotiably — assert they agree:

import time

n, nsteps = 500, 5000
u0 = np.zeros(n); u0[0] = 100.0

t0 = time.perf_counter()
a = simulate_py(u0.copy(), 0.25, nsteps)
t1 = time.perf_counter()
b = difflib.diffuse.simulate(u0.copy(), 0.25, nsteps)
t2 = time.perf_counter()

assert np.allclose(a, b)            # same answer, or the speedup is a lie
print(f"python : {t1-t0:.3f} s")
print(f"fortran: {t2-t1:.4f} s")

# Expected output (ILLUSTRATIVE — your numbers will differ):
# python :  ~3.0  s
# fortran:  ~0.005 s

Treat the times as placeholders; the order of magnitude — a compiled double loop leaving an interpreted one far behind — is the robust part, and the point of §15.5 is that you now measure it yourself.

For correctness beyond one step, use physics you know. A 1-D rod with a hot left end and cold right end, left to relax forever, settles into a linear temperature profile — the steady state of the diffusion equation with fixed ends. For n = 5 with ends 100 and 0, that analytical steady state is $[100, 75, 50, 25, 0]$. Run the simulation for many steps and the interior should approach those values:

Cell (Python index) 1 2 3
Analytical steady state 75 50 25
Long run (many steps) → 75 → 50 → 25

That the ported kernel converges to the value theory predicts — not merely that it is fast — is what lets the researcher trust it. Speed you can feel; correctness you must check against something you already know.

Discussion Questions

  1. The researcher could have kept the Python time loop and only vectorized the space sweep with NumPy. Compare that path to the full Fortran port on three axes: speed, code complexity, and how many times you cross the Python↔Fortran boundary. When would you choose each?
  2. Phase 4 got to skip the order='F' issue entirely because the array is one-dimensional. Rewrite the problem as a 2-D plate and explain exactly where the layout question would re-enter.
  3. The assert np.allclose(a, b) line uses a tolerance, not exact equality. Why is that the right choice when comparing a Python and a Fortran floating-point computation, even of the "same" formula? (You met the reason in Chapter 20's territory — floating-point is not associative.)

Your Turn: Extensions

  • Option A. Take the vectorized-NumPy middle path: keep the Python time loop but replace the inner Python loop with one NumPy slice expression. Benchmark all three (pure Python, NumPy-vectorized, Fortran). Where does NumPy land relative to the two extremes, and does the answer change with n?
  • Option B. Profile the original notebook with cProfile (or %prun in Jupyter) and confirm, with real numbers this time, that the inner loop dominates. Compare what the profiler says to the back-of-the-envelope count in Phase 1. (This previews Chapter 28.)
  • Option C. Extend the wrapped kernel to also return the time history (an nsteps × n array) so the notebook can animate the rod cooling. Which intent, and which memory-layout choice, does that output array need?

Key Takeaways

  • The two-language workflow is a diagnosis, not a rewrite: profile (or reason) to find the one loop that holds the runtime, and move exactly that.
  • A loop-carried dependency — each step needing the last step's result — is the signature of a kernel NumPy cannot vectorize away, and the clearest signal to reach for Fortran.
  • Pushing the entire iteration into Fortran, so Python calls across the boundary once, makes the crossing cost vanish — the opposite failure mode from calling a tiny kernel millions of times.
  • A port is not done when it is fast; it is done when it is fast and verified — one hand-checked step, and convergence to a steady state you can derive independently.