Profiling with cProfile and timeit: find where a script spends its time (Python)

The problem

Your analysis gives the right answer in about 3 s per run on the machine that built this page, and the parameter sweep you need has a thousand runs: most of an hour. Before you rewrite anything, you want to know which function takes the time, and after the fix you want proof that it helped and changed nothing else. Profiling answers the first question: cProfile runs your script and records the time spent in each function. timeit answers the second. The example is 2,000 random walkers diffusing in a box and their mean squared displacement, how far they have spread on average. Swap in your own script, with its work split into functions, as the next section shows.

The code

This file stands in for your script: %%writefile walkers.py on the first line of a cell saves the cell to that file instead of running it. cProfile reports time per function, so code at the top level of notebook cells shows up as one <module> row. Copy such cells into a .py file and wrap each stage in a def that takes inputs and returns results, as step, reflect, and msd do. When you write the fix, keep the slow function next to it, as msd stays next to msd_fast: the comparison at the end needs both. simulate takes either one as an argument, and a seeded generator gives both the same paths.

%%writefile walkers.py
import sys
import numpy as np

def step(pos, size, rng):
    return pos + size * rng.standard_normal(pos.shape)

def reflect(pos, box):
    return box - np.abs(box - np.abs(pos))    # bounce off the walls at 0 and at box

def msd(pos, start):
    total = 0.0
    for i in range(len(pos)):                 # one walker at a time
        dx = pos[i, 0] - start[i, 0]
        dy = pos[i, 1] - start[i, 1]
        total += dx * dx + dy * dy
    return total / len(pos)
def msd_fast(pos, start):                     # the fix: the loop as one NumPy line
    return np.mean(np.sum((pos - start) ** 2, axis=1))

def simulate(msd=msd, n_walkers=2_000, n_steps=2_000, size=0.01, box=1.0, seed=1):
    rng = np.random.default_rng(seed)
    pos = start = box * rng.random((n_walkers, 2))
    curve = np.empty(n_steps)
    for t in range(n_steps):
        pos = reflect(step(pos, size, rng), box)
        curve[t] = msd(pos, start)            # the mean squared displacement after each step
    return curve

if __name__ == "__main__":  # true when run as a program, false on import
    simulate(msd_fast if "--fast" in sys.argv else msd)  # "python walkers.py --fast" runs the fix

Writing walkers.py

Inside a Jupyter kernel, cProfile also records the kernel's own background work, and on Python 3.12 with ipykernel 7.4 the table came out wrong: a function listed as taking less time than a function it calls. So the cell sends python -m cProfile -o before.prof walkers.py with subprocess.run, the command you would type in a terminal: cProfile runs the script as a program of its own and writes the profile to before.prof, which pstats prints as a table. The --fast flag switches the script to the fix, and your script goes on the line marked # <- your script here. To measure again on this page, delete before.prof, after.prof, and timings.json.

import json, os, pathlib, pstats, subprocess, sys, timeit
import numpy as np

# ---- profile the script in its own process, as it is and with the fix
for name, flags in [("before", []), ("after", ["--fast"])]:    # *flags below adds "--fast" or nothing
    if not os.path.exists(f"{name}.prof"):  # measures once for this page; drop this "if" in your copy
        subprocess.run([sys.executable, "-m", "cProfile", "-o", f"{name}.prof",  # this notebook's Python
                        "walkers.py", *flags], check=True)     # <- your script here

# ---- read the profile
pstats.Stats("before.prof").strip_dirs().sort_stats("cumulative").print_stats(8)   # the top 8 rows
pstats.Stats("after.prof").strip_dirs().sort_stats("tottime").print_stats(4)       # the top 4 rows
pstats.Stats("before.prof").strip_dirs().print_stats(r"\(reflect\)")   # one row by name: a pattern for the last column

# ---- time the slow function alone, then the whole run
from walkers import msd, msd_fast, simulate    # imported once per session: restart the kernel after editing walkers.py
pos, start = np.random.default_rng(0).random((2, 2_000, 2))    # two arrays of 2,000 points, the real size
def best(func, number, repeat):                                 # seconds per call
    return min(timeit.repeat(func, number=number, repeat=repeat)) / number
if not os.path.exists("timings.json"):  # measures once for this page; drop this "if" in your copy
    t = {"msd": best(lambda: msd(pos, start), 500, 5), "msd_fast": best(lambda: msd_fast(pos, start), 10_000, 5),
         "run": best(lambda: simulate(msd), 1, 3), "run_fast": best(lambda: simulate(msd_fast), 1, 3)}
    pathlib.Path("timings.json").write_text(json.dumps(t))
t = json.loads(pathlib.Path("timings.json").read_text())

# ---- did the fix help?
print(f"msd alone:  {1e3 * t['msd']:.2f} ms -> {1e3 * t['msd_fast']:.3f} ms, {t['msd'] / t['msd_fast']:.0f} times faster")
print(f"one run:    {t['run']:.2f} s -> {t['run_fast']:.2f} s")
print(f"1,000 runs: {t['run'] * 1000 / 60:.0f} min -> {t['run_fast'] * 1000 / 60:.1f} min")
print(f"largest difference between the two MSD curves: {np.abs(simulate(msd) - simulate(msd_fast)).max():.1e}")

Thu Oct  8 12:16:34 2026    before.prof

         100076 function calls (98221 primitive calls) in 3.173 seconds

   Ordered by: cumulative time
   List reduced from 681 to 8 due to restriction <8>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    100/1    0.001    0.000    3.173    3.173 {built-in method builtins.exec}
        1    0.000    0.000    3.173    3.173 walkers.py:1(<module>)
        1    0.005    0.005    3.072    3.072 walkers.py:20(simulate)
     2000    2.879    0.001    2.880    0.001 walkers.py:10(msd)
        9    0.001    0.000    0.244    0.027 __init__.py:1(<module>)
     2000    0.149    0.000    0.149    0.000 walkers.py:4(step)
    125/2    0.001    0.000    0.121    0.061 <frozen importlib._bootstrap>:1349(_find_and_load)
    125/2    0.001    0.000    0.121    0.060 <frozen importlib._bootstrap>:1304(_find_and_load_unlocked)


Thu Oct  8 12:16:34 2026    after.prof

         130076 function calls (128221 primitive calls) in 0.364 seconds

   Ordered by: internal time
   List reduced from 691 to 4 due to restriction <4>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     2000    0.133    0.000    0.133    0.000 walkers.py:4(step)
     4000    0.075    0.000    0.075    0.000 {method 'reduce' of 'numpy.ufunc' objects}
     2000    0.015    0.000    0.015    0.000 walkers.py:7(reflect)
     2000    0.014    0.000    0.114    0.000 walkers.py:17(msd_fast)


Thu Oct  8 12:16:34 2026    before.prof

         100076 function calls (98221 primitive calls) in 3.173 seconds

   Random listing order was used
   List reduced from 681 to 1 due to restriction <'\\(reflect\\)'>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     2000    0.018    0.000    0.018    0.000 walkers.py:7(reflect)


msd alone:  1.29 ms -> 0.045 ms, 28 times faster
one run:    2.99 s -> 0.24 s
1,000 runs: 50 min -> 4.1 min
largest difference between the two MSD curves: 8.9e-16

The knobs

timeit.repeat(func, number=500, repeat=5) hands back five numbers, each the seconds that 500 back-to-back calls of func took, so one call costs a total divided by number. Choose number so that one measurement lasts at least 0.2 s, the mark timeit.Timer.autorange() aims for: 500 calls of msd and 10,000 of msd_fast pass it with room to spare. Keep the smallest of the five: the code did the same work in each, so whatever the others took on top went to the rest of the machine. Time on inputs of the real size: a loop and an array operation grow differently, so a small test case gives the wrong ratio. In a notebook, %timeit msd(pos, start) does the job in one line and reports a mean and a standard deviation, so compare its numbers only with each other.

ncalls counts how often a function ran. A pair such as 125/2 belongs to a function that was entered again before an earlier call had returned, here the import machinery loading one module in the middle of loading another: 125 calls in all, 2 of them not nested inside another call of the same function, which pstats calls primitive. tottime is the time in the function's own lines, cumtime adds everything it calls, each percall is the column to its left divided by the calls, and the last column names file, line, and function, as in walkers.py:10(msd). strip_dirs() shortens paths to file names, so files with the same name share a row: the __init__.py row with 9 calls adds up nine packages. The table has no percent column, so divide a function's cumtime by the total in the header line, and msd comes to about nine tenths of the run.

sort_stats("cumulative") lists the call chain first: exec, the script's <module>, simulate. Read down past these wrappers to the first function of your own whose cumtime is close to its tottime, since that one does the work itself. Once a fix is in, "tottime" ranks functions by their own work: after the fix step leads, then NumPy's reduce, the sum and the mean inside msd_fast. The fix is the loop rewritten as array arithmetic, the subject of Vectorizing loops with NumPy, and the two MSD curves differ by 8.9e-16, rounding from adding the same numbers in another order. For a quick look without a file, python -m cProfile -s cumulative walkers.py prints the table at once.

Pitfalls

Timing one run, imports included. The shell's stopwatch, time python walkers.py in a terminal, also counts the start of Python and the import of NumPy, and time.time() around a single call in a session counts whatever else the machine was doing at that moment. The table shows the import as the _find_and_load rows, Python's import machinery: a few percent of this run and several times what reflect costs in all 2,000 steps. A sweep pays it once, not a thousand times. Time the function after the imports, with timeit and its repeats.

Profiling overhead on small functions. cProfile adds a fixed cost to every function call it records, so a small function called a million times looks heavier in the table than it is. Here the header counts about 100,000 calls for the whole run, too few for that cost to matter next to the loop in msd. Before you rewrite a suspect, confirm it with timeit on its own, as the cell does for msd.

Optimizing what takes 1 % of the time. reflect does not reach the eight rows of the first table, so the cell asks for it by name: print_stats also takes a regular expression in place of a row count and prints only the rows whose last column matches. That lookup is the third block of output, the one reflect row of the first profile under that profile's header ("Random listing order" means unsorted), and its cumtime is under 1 % of the total: making it infinitely fast saves under 1 %. Rewriting the random numbers in step, the usual suspect, saves at most their 5 % or so. Start with the top row of your own code, and stop when the top rows are library calls that already work on whole arrays, as step and NumPy's reduce are after the fix. Past that point a faster line will not help; a different algorithm might.