Python for scientific computing

The processor only understands machine code

A processor executes machine code: sequences of very simple binary instructions such as “load the 8 bytes at this memory address into register 3”, “add register 3 to register 5” or “if the last result was zero, continue at that address”. The chapter on computer architecture shows a toy processor doing exactly this. Every program that runs on a computer, whether it was written in Python, Fortran or JavaScript, must in the end be turned into such instructions, because nothing else can drive the hardware.

There are two broad strategies for getting from a human-readable language to machine code. A compiler translates the whole program into machine code before it runs. Fortran and C work this way: you run gfortran md.f90 -o md, the compiler produces a file of machine code called an executable, and when you run ./md the processor executes that file directly. Compilation takes time, but the result runs fast, because the compiler has studied the whole program and chosen efficient instructions for it.

An interpreter takes the other approach. It is itself a program, already compiled to machine code, that reads your program and carries out what it says, step by step, while it runs. Nothing in your program is ever turned into machine code of its own; instead the interpreter’s machine code examines your instructions one at a time and performs them. This is slower, since the interpreter has to work out what each step means every time it reaches it, but it is enormously flexible: you can type a line and see the result immediately, change a program while it runs, and inspect anything at any moment.

Python, as it is normally used, sits between the two. The program that runs Python code is called CPython, because it is written in C. It is the reference implementation, distributed from python.org and by Anaconda, Homebrew and every Linux distribution. CPython first compiles your program, not into machine code but into a compact intermediate form called bytecode, and then interprets the bytecode with a loop written in C called the virtual machine. The rest of this chapter follows a small program through each stage.

Other implementations of the same language exist. PyPy compiles frequently executed parts of a program into machine code while it runs (just-in-time compilation) and is often several times faster on pure-Python loops. The Python inside this page is Pyodide: CPython itself, compiled to WebAssembly, the portable machine code that web browsers can run. When you press a run button on this site, it is the genuine CPython 3.12 interpreter that executes your code, just running inside the browser rather than directly on your operating system.

Starting the interpreter

When you type python3 analyse.py in a terminal, the shell asks the operating system to start a new process running the CPython executable, a file of machine code typically found at /usr/bin/python3 or inside a conda environment, and passes it the argument analyse.py (the chapter on how bash works describes this in detail). CPython initialises itself: it sets up its memory allocator, creates the built-in objects such as int, print and None, and imports a few essential modules. Then it reads your file and the pipeline below begins.

There are other ways to give CPython code. Typing python3 with no file starts the interactive interpreter, or REPL (read–evaluate–print loop), which compiles and runs each statement as you enter it and prints the value of each expression. A Jupyter notebook sends cells of code to a long-running Python process called a kernel, which compiles and runs each cell when you execute it and keeps all names alive between cells. The cells on this site behave the same way: cells on one page share an interpreter, so a name defined in one cell can be used in the next.

From characters to tokens

To the computer, a program file is just a sequence of bytes. Python source files are encoded in UTF-8, in which every ordinary English letter, digit and symbol is one byte, the same value as in the old ASCII code, while characters such as π, Å or Chinese characters take two to four bytes. Python allows such characters in names and strings, so Δx = 0.1 is a legal statement.

The first stage, the tokenizer, reads these characters and groups them into the words of the language, called tokens. In y = 2 * (x + 4) there are nine: the name y, the operator =, the number 2, the operator *, an opening bracket, the name x, the operator +, the number 4 and a closing bracket. Spaces between tokens are discarded. Line structure is not: the tokenizer emits a NEWLINE token at the end of every logical line, and an INDENT or DEDENT token whenever the indentation increases or decreases. This is how Python’s indentation becomes part of the grammar. Where C and Fortran use braces or end do to close a block, Python uses a DEDENT token generated from the whitespace.

The figure below runs the real CPython pipeline on any code you type. Press the button, then click through the stages.

From tokens to a syntax tree

The second stage, the parser, checks that the sequence of tokens obeys Python’s grammar and builds an abstract syntax tree (AST) that records the structure of the program. Look at the syntax tree in the figure for y = 2 * (x + 4). At the top is an Assign node with a target, the name y in “store” context, and a value. The value is a BinOp node for the multiplication, whose left operand is the constant 2 and whose right operand is another BinOp for the addition of x and 4. The brackets have disappeared, because the tree’s shape already says that the addition happens first. Remove the brackets and compile again: now the multiplication 2 * x is nested inside the addition, because * binds more tightly than +. Operator precedence is not a separate rule applied at run time; it is decided once, by the shape of the tree.

If the tokens do not form a valid program, parsing fails and Python reports a SyntaxError before running anything at all. Type y = 2 * (x + 4 with the closing bracket missing and compile: there is no partial output, because no bytecode was ever produced. This is different from an error such as dividing by zero, which can only be discovered while the program runs. Distinguishing errors found while reading a program from errors found while running it is a useful habit when debugging.

From the tree to bytecode

The third stage, the compiler, walks the syntax tree and emits bytecode. Bytecode is a list of simple instructions, each two bytes long (some followed by a few bytes of cache space), for an imaginary machine that does not exist in hardware: the CPython virtual machine. Open the bytecode stage of the figure. The assignment x = 3 becomes two instructions, LOAD_CONST 3 and STORE_NAME x; the expression 2 * (x + 4) becomes “load 2, load x, load 4, add, multiply” in exactly the order the tree prescribes, so that the addition happens before the multiplication.

The compiled result is a code object, containing the bytecode, the table of constants it uses (here 3, 2, 4 and None), the names it refers to, and a table mapping each instruction back to a line of the source, which is how error messages can tell you the line where something went wrong. Every function has a code object of its own, compiled once when the module is compiled.

Compiling is fast, but not free, so CPython caches the bytecode of imported modules. The first time you import mymodule, CPython writes the compiled code to a file such as __pycache__/mymodule.cpython-312.pyc, and later imports load the .pyc file directly if the source has not changed. That is the origin of the __pycache__ directories that appear next to your code. You can look at any bytecode yourself with the standard dis module:

The virtual machine

The bytecode is executed by the evaluation loop, a large function in CPython’s C source (Python/ceval.c, with the instruction bodies defined in Python/bytecodes.c since version 3.12). In essence it is a loop that fetches the next instruction, looks at its opcode, jumps to the C code that implements that opcode, runs it, and repeats: a software imitation of a processor’s fetch–decode–execute cycle.

CPython’s virtual machine is a stack machine. It has no registers. Instead, each instruction takes its inputs from the top of a value stack and leaves its result there. LOAD_CONST pushes a constant. LOAD_NAME looks up a name and pushes the object it refers to. BINARY_OP + pops two values, adds them and pushes the sum. STORE_NAME pops the top value and binds a name to it. The figure below executes real CPython 3.12 bytecode for four short programs, one instruction at a time, and explains each instruction as it goes.

Step through the four programs and some important facts about Python become visible. An if statement compiles to a comparison followed by a conditional jump, POP_JUMP_IF_FALSE: a decision is just a jump to a different instruction. A for loop compiles to GET_ITER, which creates an iterator object, then a FOR_ITER instruction that asks the iterator for the next value, the loop body, and a JUMP_BACKWARD that goes round again until the iterator is exhausted. range(4) produces numbers on demand rather than building a list of them.

The function-call program shows frames. A function definition, def square(a), is not a declaration but an action performed at run time: MAKE_FUNCTION wraps the precompiled code of square in a new function object and STORE_NAME binds the name square to it. When square(5) is called, CALL creates a new frame, a block of memory holding the function’s own value stack, its local variables and a pointer to where execution is in its bytecode. Execution continues in the new frame until RETURN_VALUE destroys it and hands the result back to the caller. The chain of frames waiting for each other is the call stack, which is what a traceback lists when an error occurs. Local variables are read with LOAD_FAST, which fetches them from a numbered slot in the frame, rather than LOAD_NAME, which searches dictionaries; this is why code inside functions runs noticeably faster than the same code at the top level of a script.

Why Python is slower than C, and what has been done about it

Each bytecode instruction is carried out by C code that typically executes dozens to hundreds of machine instructions. The work is not only the arithmetic itself. To execute BINARY_OP + the interpreter must fetch and decode the instruction, check the types of both operands (they could be integers, floats, strings, lists or objects of a class you defined), find the right addition routine for those types, check whether the result fits, create a new object to hold it, and adjust reference counts. A compiled Fortran loop adding two arrays of numbers spends perhaps one or two machine instructions per element, because the compiler knew in advance that every element is a 64-bit float. The flexibility of Python, where any name may refer to any type of object at any moment, is paid for at every step.

This is not a reason to avoid Python; it is a reason to use it well. The usual pattern in scientific computing is to write the overall logic in Python and to hand the heavy numerical loops to compiled code: a NumPy operation on an array of a million numbers is a single Python-level instruction that runs a compiled C loop internally, so the interpreter’s overhead is paid once rather than a million times (the chapter on NumPy explains how). Machine-learning frameworks such as PyTorch work the same way, with Python describing the computation and compiled kernels on the GPU doing it.

CPython itself has also become much faster. Since version 3.11 it has been a specialising adaptive interpreter (PEP 659): after an instruction has run a few times, CPython replaces it with a specialised version for the types it has actually seen, such as an addition that assumes two floats, guarded by a quick check. Version 3.13 added an experimental just-in-time compiler and an optional build without the global interpreter lock, which is discussed with parallelism in the chapter on performance.

Every value is an object in memory

One more part of the machinery shapes everything else about the language. In CPython, every value, including every integer, float and string, is an object: a block of memory on the heap that begins with a header containing a reference count and a pointer to the object’s type. The integer 1 is not a bare 8-byte number as it would be in Fortran; on a 64-bit build it is a 28-byte structure. Names, list slots and dictionary entries do not contain values; they contain pointers to objects.

The reference count records how many references point to the object. Binding a new name to it increments the count; deleting the name, rebinding it or leaving a function decrements it; and when the count reaches zero the memory is freed immediately. A separate garbage collector occasionally searches for groups of objects that refer to each other in a cycle and are unreachable from anywhere else, which reference counting alone cannot free.

This model, in which names are labels attached to objects rather than boxes containing values, explains behaviour that surprises almost every newcomer, such as two names that change together when a list is modified through one of them. The next chapter makes it visible with a tool that draws every name and object in memory as a program runs.

References and further reading

Every value is an object

The previous chapter ended with a fact about CPython’s machinery: every value lives in memory as an object. This chapter builds the whole of Python’s data model on that fact, because nearly everything that confuses newcomers about the language, and a good deal that confuses experienced programmers, follows from it.

Every object has three properties. Its identity is a number that is unique among the objects alive at the same time; in CPython it is the object’s memory address, returned by the built-in function id(). Its type determines what the object is and what can be done with it: the integer 300 has type int, the text "argon" has type str, a list has type list, and type(x) returns the type of any object. Its value is the data it carries. Identity and type never change during an object’s lifetime. Whether the value can change depends on the type: some objects, such as lists, can be modified in place, and others, such as integers and strings, cannot.

The type decides which operations are allowed and what they mean. + adds two integers, concatenates two strings and joins two lists, because each type defines its own addition; 3 + "4" fails with a TypeError because neither type knows how to add itself to the other. Types are themselves objects, of type type, and they can be created by your own code with the class statement (the chapter on classes).

Names are references

An assignment such as x = 300 does two things. It evaluates the right-hand side, creating (or finding) an object, and then it binds the name x to that object. The name is not a box that contains 300; it is a label that refers to an object elsewhere in memory. A name has no type of its own and can be rebound at any time to an object of any type.

The consequences become clear only when they can be seen. The figure below runs a program with real Python, pausing before each line, and draws every name on the left and every object on the right, with an arrow for each reference.

Step through the first example. After y = x, there are two names and one object: the assignment copied the reference, not the integer. Then x = x + 1 computes a new integer object, 301, and rebinds x to it; y still refers to 300. Nothing was modified. Integers cannot be changed, so “changing x” always means pointing x at a different object.

The second example shows the same rule with a list, and the result is the one that surprises people. b = a binds a second name to the same list. b.append(4) does not rebind anything: it modifies the list object in place, so the change is visible through both names. Two names referring to one mutable object are called aliases. a.copy() creates a new list object with the same elements, so appending to c leaves a alone.

The last example shows the most common bug that aliasing causes. [row, row] is a list containing two references to one list, so setting grid[0][0] changes “both rows”. A list comprehension, [[0, 0] for _ in range(2)], evaluates [0, 0] afresh on each pass and creates two independent rows. The same trap catches people who write [[0] * n] * m to make a matrix; a NumPy array, np.zeros((m, n)), avoids it entirely.

Function arguments follow the same rule. Calling f(a) binds the function’s parameter to the same object that a refers to. If the function modifies that object in place, the caller sees the change; if it rebinds its parameter, the caller does not. Python’s argument passing is therefore neither “by value” nor “by reference” in the C++ sense; it is sometimes called call by sharing, and it is just assignment.

Numbers

Python has three built-in numeric types. Integers (int) have unlimited precision. Where a Fortran integer(8) overflows above $2^{63} - 1 \approx 9.2 \times 10^{18}$, a Python integer simply grows: 2**100 is 1267650600228229401496703205376, exactly. CPython stores large integers as arrays of 30-bit “digits”, so their size in memory grows with the number of digits, and arithmetic on them is correspondingly slower. This makes Python convenient for exact combinatorics, such as counting the microstates of an Einstein solid, but it is not how numerical arrays are stored.

Floats (float) are IEEE 754 double-precision numbers, exactly the 64-bit format described in the chapter on numbers in a computer, with about 16 significant decimal digits and a range up to about $1.8 \times 10^{308}$. Everything said there applies: 0.1 + 0.2 is 0.30000000000000004, because none of the three numbers is exactly representable in binary, so floats should be compared with a tolerance using math.isclose rather than with ==. Overflow gives inf, and undefined operations such as inf - inf give nan, which is not equal even to itself. Complex numbers (complex) are pairs of floats written as 2 + 3j.

The arithmetic operators are +, -, *, / (true division, which always returns a float), // (floor division), % (remainder) and ** (power). Floor division rounds towards minus infinity, so -7 // 2 is -4, and the remainder takes the sign of the divisor, so that (a // b) * b + a % b == a always holds. This differs from Fortran and C, which truncate towards zero; it matters when wrapping coordinates into a periodic box, where Python’s x % L conveniently always lands in $[0, L)$. Mixed arithmetic promotes integers to floats and floats to complex. The math module provides sqrt, exp, log, trigonometric functions and constants such as math.pi for single numbers; for arrays, NumPy provides vectorised versions of all of them.

Booleans (bool) are the two objects True and False, produced by comparisons such as x > 0. The type is a subclass of int, so True + True == 2, which is occasionally handy for counting how many conditions hold.

Strings

A string (str) is an immutable sequence of Unicode characters. Strings are written in single or double quotes, or in triple quotes for text that spans several lines. Individual characters are reached by indexing from zero, s[0], and from the end with negative indices, s[-1]. A slice s[start:stop] extracts the characters from start up to but not including stop, so "argon"[1:3] is "rg". The half-open convention, inherited by lists and NumPy arrays, has two advantages: s[:k] and s[k:] split a sequence into two parts with no overlap, and len(s[a:b]) == b - a.

Strings come with many methods: .upper(), .strip() to remove surrounding whitespace, .split() to break a line of a data file into fields, .replace(), .startswith(), and ", ".join(items) to glue a list of strings together. Because strings are immutable, every method returns a new string; s.upper() does not change s. The third example in the figure shows this: s = s + 'gon' creates a new string and rebinds s, while t still refers to the original 'Ar'.

The most useful string tool in scientific work is the f-string, which inserts values into text with a format specification:

The part after the colon controls the format: .4e means scientific notation with four digits after the point, 10.3f a fixed-point number three decimals wide padded to ten characters, >6d a right-aligned integer six characters wide. Writing tables of results this way produces files that line up and can be read back reliably.

None and truth

None is the single object of type NoneType. It represents the absence of a value: a function without a return statement returns None, and None is the conventional default for an optional argument. Test for it with x is None.

In a condition, every object counts as true or false. The false values are False, None, zero of any numeric type, and empty containers: "", [], (), {} and set(). Everything else is true, including the string "False" and a list containing a zero, [0]. This allows idioms such as if atoms: meaning “if the list of atoms is not empty”. It also allows bugs: if x: is false when x is 0.0, which may be a perfectly valid coordinate, so test if x is not None: when that is what you mean.

Mutability, identity and equality

Types divide into immutable ones, whose objects never change after creation (int, float, complex, bool, str, tuple, frozenset, bytes), and mutable ones, whose objects can be modified in place (list, dict, set, bytearray, NumPy arrays and most user-defined classes). For an immutable object, aliasing is harmless, since nothing can change it behind your back. For a mutable object, every alias is a window onto the same data.

Two different questions can be asked of two objects. a == b asks whether their values are equal, which each type defines for itself. a is b asks whether they are the same object, which compares identities. Two separately created lists [1, 2] and [1, 2] are equal but not identical. Use is only for singletons such as None, True and False. For numbers and strings its result depends on implementation details: CPython keeps a single shared object for each small integer from $-5$ to 256 and reuses identical string constants, so 7 is 7 happens to be true while two separately computed 1000s may or may not be the same object. Never rely on it.

Mutability also determines what can be a dictionary key or set element. A dictionary finds a key by its hash, a number computed from the value, and that number must not change while the key is stored; so keys must be immutable. Tuples of numbers make good keys, for example a pair of atom indices, while lists cannot be used. The chapter on collections shows how hashing works.

Copying has two levels. A shallow copy, a.copy(), list(a) or a[:], creates a new outer container whose slots refer to the same element objects. If the elements are themselves mutable, such as the rows of a list of lists, the copy and the original still share them. A deep copy, copy.deepcopy(a), recursively copies everything reachable. In numerical work, NumPy’s arr.copy() copies the whole numerical buffer, which is usually what you want.

Memory management

CPython frees memory automatically using reference counting. Every object records how many references point to it. Binding a name, storing the object in a list or passing it to a function increments the count; rebinding the name, deleting it with del, removing it from the list or returning from the function decrements it. When the count reaches zero, the object is destroyed immediately and its memory returned. In the figure, every object that loses its last arrow disappears from the picture at the next step, because nothing can reach it any more.

Reference counting cannot free a group of objects that refer to each other in a cycle, such as two objects each holding a reference to the other, because their counts never fall to zero. CPython therefore also runs a cyclic garbage collector from time to time, which finds groups of objects reachable only from each other and frees them. In everyday code you never call it; in a long simulation that creates large cyclic structures, import gc; gc.collect() can be used to force a collection.

The object model has a memory cost that matters for scientific data. On a 64-bit build, each integer object occupies 28 bytes and each float 24, and a list adds an 8-byte pointer per element on top. A list of a million floats therefore uses about 32 MB, whereas a NumPy array of a million float64 values stores the bare numbers contiguously in exactly 8 MB, with no per-element objects at all. For large numerical data, NumPy is the right container, and the chapter on NumPy arrays explains its memory layout.

Dynamic typing and type hints

Python is dynamically typed: types belong to objects, not to names, and are checked only when an operation is attempted. x = 1 and then x = "one" is perfectly legal. A mistake such as adding a number to a string is discovered only when that line runs, perhaps hours into a calculation. It is also strongly typed: it will not silently convert "3" into the number 3, unlike some languages.

Since Python 3.5, code can carry optional type hints, which state what types are expected without changing how the program runs:

def kinetic_energy(mass: float, velocity: list[float]) -> float:
    return 0.5 * mass * sum(v * v for v in velocity)

The interpreter ignores these annotations, but editors use them to offer completion and to warn about mistakes, and static checkers such as mypy can analyse a whole code base for type errors before it runs. For scientific code that will be shared or maintained, type hints on function signatures are a cheap form of documentation that tools can verify.

References and further reading

Statements, blocks and indentation

A Python program is a sequence of statements executed one after another, top to bottom. Assignments, function calls and import are simple statements that occupy a line. Compound statements, such as if, while, for, def and try, control the execution of other statements: they consist of a header line ending in a colon, followed by an indented block of statements that belong to it.

if temperature > melting_point:
    phase = "liquid"          # inside the if block
    print("melted")           # also inside
print("done")                 # after the block: always runs

Indentation is not decoration. As the chapter on how Python runs showed, the tokenizer turns increases and decreases of indentation into INDENT and DEDENT tokens, and the parser uses them exactly as C uses braces or Fortran uses end if. All lines of a block must be indented by the same amount; four spaces per level is the universal convention (PEP 8). Mixing tabs and spaces is a classic source of IndentationError, which is why editors are usually set to insert spaces when the tab key is pressed.

Underneath, every control structure in this chapter compiles to the same two primitives seen in the virtual-machine figure of that chapter: conditional jumps, which skip to another instruction if a value is false, and unconditional jumps, which go back to repeat a block. Everything else is convenience built on top.

Conditions and Boolean logic

A condition is any expression whose truth value can be tested. Most conditions are comparisons: == (equal), != (not equal), <, <=, >, >=, together with in (membership, as in "Fe" in elements) and is (identity). Comparisons can be chained in the mathematical way: 0 <= x < L means 0 <= x and x < L, which is exactly how one checks that a coordinate lies inside a periodic box.

Conditions are combined with and, or and not. The first two use short-circuit evaluation: in a and b, if a is false the result is already known and b is never evaluated; in a or b, if a is true, b is skipped. This is not merely an optimisation, it is relied upon for safety, as in if n > 0 and total / n > threshold:, where the division is attempted only when it is safe. (The same idea in the shell, make && ./run, is described in the chapter on bash.) Strictly, and and or return one of their operands rather than True or False: x or default gives x if it is truthy and default otherwise, an idiom you will meet in other people’s code.

The if statement chooses between blocks:

if energy < e_min:
    status = "below range"
elif energy > e_max:
    status = "above range"
else:
    status = "ok"

The conditions are tested in order, and only the block of the first true one runs; elif and else are optional. For a choice between two values there is a one-line conditional expression, status = "ok" if e_min <= energy <= e_max else "out of range". Python 3.10 added the match statement for structural pattern matching, useful for dispatching on the shape of data, but a chain of if/elif covers most scientific needs.

Repeating while a condition holds

A while loop repeats its block as long as its condition is true, testing the condition before each pass. It is the natural loop when the number of repetitions is not known in advance, which in numerical work usually means iterating until convergence. Newton’s method for $\sqrt a$ improves a guess $x$ with the rule $x \leftarrow \tfrac12(x + a/x)$ (the chapter on derivatives derives it), and stops when $x^2$ is close enough to $a$:

Newton’s method reaches full double precision in five iterations, roughly doubling the number of correct digits each time; bisection gains one binary digit per halving and needs 34 halvings to shrink the interval below $10^{-10}$. Step through the convergence example in the figure to see the same loop line by line.

Two precautions apply to every while loop. The condition must eventually become false, or the loop runs forever; in an iterative solver that may fail to converge, add a maximum number of iterations and raise an error if it is reached. And the tolerance must be achievable: asking for $|x^2 - a| < 10^{-20}$ can never succeed in double precision, because the rounding error of $x^2$ itself is larger than that. x += 1 is shorthand for x = x + 1; the same augmented form exists for every operator (-=, *=, /=).

For loops and the iterator protocol

The for loop is Python’s main loop, and it works differently from the counting loops of C and Fortran. It does not count; it iterates: it takes any iterable object, a collection or a stream of values, and runs its block once for each value, binding the loop variable to each in turn.

for element in ["H", "He", "Li"]:      # a list
    print(element)
for line in open("traj.xyz"):          # the lines of a file, read one at a time
    ...
for k in range(10):                    # the integers 0, 1, ..., 9
    ...

The machinery behind this is the iterator protocol, and it is worth understanding because it explains why so many things can be looped over. A for loop first calls iter() on the object, which returns an iterator: an object that remembers a position and has a __next__ method. The loop then calls next() on the iterator repeatedly, binding each result to the loop variable, until next() signals that there are no more values by raising the StopIteration exception, which the loop catches quietly. In bytecode, these are exactly the GET_ITER and FOR_ITER instructions seen in the virtual-machine figure.

range(start, stop, step) produces integers on demand, half-open like slices: range(10) gives 0 to 9, range(2, 10, 2) gives 2, 4, 6, 8. It does not build a list, so range(10**9) takes no memory. Two built-ins combine iterables in ways that remove most of the need for indices. enumerate(seq) yields pairs (index, item), and zip(a, b) walks two sequences in step, yielding pairs (a[i], b[i]) and stopping at the shorter one:

Writing for i in range(len(elements)): and then elements[i] works, but it is slower, longer and invites off-by-one errors; iterating directly over the items is the idiomatic style. Dictionaries are iterated with .items() for key–value pairs, .keys() or .values().

Changing the course of a loop

Three statements alter a loop from inside. break leaves the innermost loop immediately. continue abandons the rest of the current pass and goes on to the next one, which is handy for skipping comment lines when reading a file:

for line in open("data.txt"):
    if line.startswith("#"):
        continue                 # skip comments
    x, y = map(float, line.split())

The third is unusual: a loop may have an else clause, which runs if the loop finished normally, without a break. It reads awkwardly but expresses “searched everything and found nothing” exactly. In the prime-number example in the figure, the inner loop looks for a divisor of n; if it finds one it breaks, and if it finishes without breaking, the else block records n as prime. Step through it and notice how many times the inner loop runs for each n.

Nested loops multiply their costs: a double loop over all pairs of $N$ atoms executes its body $N^2$ times, or $N(N - 1)/2$ if each pair is visited once with for j in range(i + 1, N). In pure Python, where each pass costs many bytecode instructions, a double loop over 10 000 atoms (50 million pairs) takes tens of seconds; the same calculation done with NumPy arrays, or with a cell list that only visits nearby pairs (the chapter on neighbour lists), takes a fraction of a second.

Comprehensions and generators

Building a list with a loop is so common that Python has a compact expression for it, the list comprehension:

squares = [k * k for k in range(6)]                 # [0, 1, 4, 9, 16, 25]
even_squares = [k * k for k in range(6) if k % 2 == 0]   # [0, 4, 16]
distances = [abs(x - x0) for x in positions]

The comprehension reads like set-builder notation in mathematics, $\{k^2 : k \in \{0, \dots, 5\}, k \text{ even}\}$, and runs somewhat faster than the equivalent loop with append. The same syntax with braces builds a set or a dictionary: {el: m for el, m in zip(elements, masses)}.

Replacing the brackets with parentheses gives a generator expression, which produces the values one at a time on demand instead of storing them all. sum(k * k for k in range(1_000_000)) adds a million squares without ever holding a million-element list in memory. More generally, a function containing the keyword yield is a generator function: each yield hands one value to the caller and pauses the function, with all its local variables intact, until the next value is requested. Generators are the natural way to write a reader that streams frames from a trajectory file too large to load at once:

def frames(path):
    with open(path) as f:
        while True:
            header = f.readline()
            if not header:
                return                       # end of file: stop the iteration
            n = int(header)
            f.readline()                     # comment line
            yield [f.readline().split() for _ in range(n)]

for atoms in frames("traj.xyz"):             # one frame in memory at a time
    ...

Exceptions

Not every change of course is planned. When an operation cannot be carried out, such as opening a file that does not exist, converting "abc" to a float or dividing by zero, Python raises an exception. Raising abandons the current block and every enclosing block and function call, unwinding the call stack frame by frame, until a matching handler is found. If none is found, the program stops and prints a traceback listing the frames that were active, most recent call last, which is the first thing to read when debugging.

Exceptions are handled with try:

try:
    value = float(field)
except ValueError:
    value = float("nan")          # mark unreadable entries instead of crashing
else:
    n_good += 1                   # runs only if no exception occurred
finally:
    n_seen += 1                   # runs in every case

Catch the specific exception you expect, never a bare except:, which also swallows typing mistakes and keyboard interrupts. Your own code should raise exceptions when it is given impossible input, such as raise ValueError(f"temperature must be positive, got {T}"): failing loudly at the point of the error is far better than producing a plausible-looking wrong number hours later. The common built-in exceptions are worth recognising: TypeError (wrong type of object), ValueError (right type, bad value), IndexError and KeyError (missing position or key), FileNotFoundError, ZeroDivisionError, and NameError (a name that was never bound, often a misspelling).

A complete loop: a random walk

The pieces of this chapter combine into the inner loop of a stochastic simulation. A particle on a line takes a thousand steps, each one unit left or right at random; repeating the walk many times shows that the typical distance travelled grows as the square root of the number of steps, the signature of diffusion.

Quadrupling the number of steps doubles the root-mean-square distance, as the central limit theorem predicts. The same law, measured from the mean-squared displacement of atoms, gives the diffusion coefficient in trajectory analysis. The program also illustrates the limits of pure-Python loops: it executes about 4.2 million passes of the inner loop and takes a second or two, where NumPy would generate all the steps as one array and finish almost instantly.

References and further reading

Why functions

A function packages a computation under a name so that it can be used without thinking about how it works. A function that returns the Lennard-Jones energy of two atoms can be called a million times inside a simulation loop, tested once on its own, and replaced by a faster or more accurate version without touching the code that calls it. Functions are the main tool for controlling complexity in a program, in the same way that lemmas control complexity in a proof: each one does one clearly stated job, and larger results are built from them.

A function is defined with def:

def lennard_jones(r, epsilon=1.0, sigma=1.0):
    """Pair energy of two atoms at distance r (same length units as sigma)."""
    sr6 = (sigma / r) ** 6
    return 4 * epsilon * (sr6 * sr6 - sr6)

The header names the function and its parameters; the indented body runs each time the function is called; return hands a value back to the caller and ends the call. A function that reaches the end of its body without return returns None. The string in triple quotes on the first line of the body is the docstring: it documents what the function does, and help(lennard_jones) displays it. Good docstrings state the meaning and units of each argument and of the result, which in scientific code are the things most often misunderstood.

As the chapter on how Python runs showed, def is an executable statement. The body is compiled once, together with the rest of the module, into a code object; when the def line is executed, MAKE_FUNCTION wraps that code object in a function object and binds the name to it. A function is therefore an ordinary object, like an integer or a list: it can be bound to another name, stored in a list or passed to another function, possibilities used later in this topic.

Arguments

When a function is called, its parameters are bound to the arguments supplied. Arguments can be given by position, lennard_jones(1.5), or by name, lennard_jones(r=1.5, sigma=1.2). Parameters with default values may be omitted. Keyword arguments make calls self-documenting: run(steps=10000, dt=0.002, temperature=300) can be understood without looking up the order of the parameters, and Python rejects a misspelt keyword immediately.

Two special forms collect a variable number of arguments. A parameter written *values gathers all extra positional arguments into a tuple, and **options gathers extra keyword arguments into a dictionary. Parameters placed after a bare * must be given by keyword, which is a good way to force clarity for options that are easily confused:

def mean(*values):
    return sum(values) / len(values)

def integrate(f, a, b, *, n=1000, method="simpson"):   # n and method keyword-only
    ...

Binding parameters to arguments is ordinary assignment, so the rules of objects and names apply. The function’s parameter and the caller’s variable refer to the same object. A function that appends to a list it was given modifies the caller’s list; a function that rebinds its parameter, values = values + [x], does not.

The same rule produces a notorious trap. Default values are evaluated once, when the def statement runs, not each time the function is called. A mutable default is therefore shared by every call that uses it:

def append_to(item, bucket=[]):    # one list, created when the def runs
    bucket.append(item)
    return bucket

append_to(1)    # [1]
append_to(2)    # [1, 2]  -- the same list again

The figure below shows the shared list sitting in memory between calls. The standard fix is to use None as the default and create the list inside the function: if bucket is None: bucket = [].

Scope

Every name lives in a namespace, and the region of code in which a name is visible is its scope. Names assigned inside a function are local to that function: they are created in the function’s frame when it runs and vanish when it returns, and a local name hides any name of the same spelling outside. This is what makes functions safe to combine: a variable i inside one function cannot clash with an i in another.

When Python looks up a name, it searches four scopes in order, known as the LEGB rule: the Local scope of the current function, then the scopes of any Enclosing functions (for functions defined inside other functions), then the Global scope of the module, and finally the Built-in scope, which holds names such as print, len and range. The first match wins. This is why you can accidentally hide a built-in by naming a variable sum or list, after which sum(...) no longer works in that scope.

The compiler decides at compile time which names are local: any name that is assigned anywhere in a function is local throughout that function. This leads to a surprising error:

x = 10
def g():
    print(x)     # UnboundLocalError: x is local to g, because of the next line
    x = 5

Because x is assigned in g, every use of x in g refers to the local variable, which does not yet exist when print runs. To rebind a module-level name from inside a function, declare it with global x; to rebind a name of an enclosing function, use nonlocal. Both should be rare. A function that reads and writes global state is hard to test and hard to reason about; passing values in as arguments and returning results is almost always better.

The call stack and recursion

When one function calls another, the caller’s frame stays alive, waiting, while a new frame is created for the callee. The frames form a stack: the most recent call is on top, and returning pops it off and resumes the frame below. The figure shows this directly. Step through the factorial example: factorial(4) calls factorial(3), which calls factorial(2), and so on, until five frames for the same function exist at once, each with its own value of n. Then the calls return one by one, each multiplying the value returned from below.

A function that calls itself is recursive. Recursion is natural for problems that contain smaller copies of themselves: traversing a tree, a nested data structure or a directory hierarchy, or dividing a problem in two repeatedly. Every recursion needs a base case that returns without recursing, here n == 0; without it the stack grows until Python stops it with a RecursionError, by default after 1000 nested calls. Because each call costs a frame and Python does not optimise tail calls, deep recursion is better written as a loop in Python.

The call stack is also what a traceback shows. When an exception goes unhandled, Python prints each frame that was active, from the outermost call to the one where the error occurred, with the file, line and source text of each. Reading a traceback from the bottom up tells you what failed and then, frame by frame, how the program got there.

Functions as values

Because functions are objects, they can be passed to other functions. A numerical derivative can be written once for any function:

A lambda is a small anonymous function written as an expression, lambda x: x**3, useful for one-off functions such as sort keys. Functions that take or return other functions are called higher-order functions. They are everywhere in scientific Python: scipy.optimize.minimize(f, x0) minimises any function you give it, scipy.integrate.solve_ivp(rhs, ...) integrates any system of differential equations, and machine-learning frameworks take loss functions as arguments.

A function defined inside another function can use the enclosing function’s variables, and it keeps them alive after the outer function has returned. The combination of a function and the variables it captured is a closure. In the last example of the figure, make_counter returns the inner function increment, which carries the variable count with it; each call of tick() increments that private count. Closures are a lightweight alternative to classes for attaching state to a function.

A decorator is a function that takes a function and returns a modified version, applied with the @ syntax above a def. @functools.lru_cache remembers previous results so that an expensive function is not recomputed for the same arguments; @numba.njit compiles a numerical function to machine code; @property turns a method into an attribute. You will use decorators long before you need to write one.

Modules and packages

A module is a file of Python code, and import makes its contents available. import math finds the module, executes it once (later imports reuse the already-loaded module object), and binds the name math to the module object, whose attributes are the names the module defined: math.sqrt, math.pi. Variants bind other names: from math import sqrt, pi binds just those names, and import numpy as np binds a shorter alias, a convention so universal that np, plt for matplotlib.pyplot and pd for pandas are understood by every reader. Avoid from module import *, which floods your namespace with unknown names.

Your own code becomes a module simply by saving it as a .py file. Python finds modules by searching the directories listed in sys.path: the directory of the script being run, the standard library and the site-packages directory of the active environment. A directory of modules with an __init__.py file is a package, imported as import mypackage.analysis.

When a file is run as a script, Python sets the module variable __name__ to "__main__"; when it is imported, __name__ is the module’s name. The idiom

if __name__ == "__main__":
    main()

at the bottom of a file therefore runs main() only when the file is executed directly, so the same file can serve both as a script and as a library of functions for other scripts and for tests.

Python’s standard library is large, and knowing what is in it saves writing code. Among the modules most used in scientific work are math and cmath, random, statistics, itertools (combinations, permutations and other iteration tools), collections (Counter, defaultdict, deque), pathlib (file paths), json and csv, time and datetime, subprocess (running other programs, such as a simulation engine), argparse (command-line options) and multiprocessing. Third-party packages, above all NumPy, SciPy, Matplotlib, pandas, ASE and PyTorch, are installed with pip install or conda install into an environment: an isolated directory with its own interpreter and packages, created with python -m venv or conda create. Keeping one environment per project, with the package versions recorded in a requirements.txt or environment.yml file, is what makes a computational result reproducible on another machine or a cluster (the chapter on Linux for simulation covers the cluster side).

Writing good scientific functions

A few habits make functions in scientific code much more trustworthy. Prefer pure functions, whose result depends only on their arguments and which change nothing outside themselves: they can be tested in isolation and called from anywhere, including in parallel. Keep each function to one task that can be described in a sentence. State units in the docstring, or better, choose a consistent unit system for the whole program and say so once.

Above all, test functions against cases where the answer is known. The Lennard-Jones energy at its minimum, $r = 2^{1/6}\sigma$, must be exactly $-\varepsilon$; a derivative of $\sin$ at zero must be 1; a function that computes kinetic energy must give zero for zero velocity. Such checks can be written as assert statements and collected into test files that a tool such as pytest runs automatically:

def test_lj_minimum():
    r_min = 2 ** (1 / 6)
    assert abs(lennard_jones(r_min) - (-1.0)) < 1e-12

A few minutes spent on tests like this routinely catches the sign errors and unit mistakes that would otherwise invalidate weeks of simulation.

References and further reading

Containers and the cost of operations

Most programs spend their time manipulating collections of things: the atoms of a structure, the frames of a trajectory, the lines of a file, the energies of a set of candidate materials. Python has four general-purpose built-in containers: the list, the tuple, the dictionary and the set. Each is good at some operations and poor at others, and choosing well can make the difference between a script that runs in a second and one that runs for an hour.

The language used to compare them is algorithmic complexity, written with “big-O” notation. An operation is $O(1)$, or constant time, if its cost does not depend on the size $n$ of the container: reading the 5th element of a list takes the same time whether the list has ten elements or ten million. It is $O(n)$, linear time, if the cost grows in proportion to $n$: checking whether a value occurs in an unsorted list may require looking at every element. Sorting is $O(n \log n)$. The notation ignores constant factors and describes how the cost scales, which is what matters when data grow. An $O(n)$ operation inside a loop that runs $n$ times is $O(n^2)$ in total; for $n = 10^5$, that is $10^{10}$ steps, which is why an innocent-looking if x in big_list: inside a loop can bring a program to a halt.

Lists

A list is an ordered, mutable sequence of references to objects, written in square brackets: masses = [1.008, 12.011, 15.999]. The elements may be of any types, mixed freely, since each slot holds only a pointer. Lists are indexed from zero, accept negative indices counted from the end, and support slicing exactly as strings do: masses[1:] is a new list of all but the first element, and masses[::-1] is the list reversed. Unlike strings, lists can be changed: masses[0] = 1.00794 replaces an element, masses[1:3] = [] deletes a slice, and the methods append, extend, insert, pop, remove, sort and reverse modify the list in place.

To understand what these operations cost, it helps to know how CPython stores a list. The list object holds a pointer to a separate contiguous array of slots, each slot an 8-byte pointer to an element object (on a 64-bit build). Reading or writing a[i] computes the address of slot i directly, which is why indexing is $O(1)$. Inserting at the front, a.insert(0, x), must shift every existing pointer one slot along, so it is $O(n)$; so is a.pop(0). For a queue that grows at one end and shrinks at the other, collections.deque does both ends in $O(1)$.

Appending is the interesting case. If the slot array were exactly as long as the list, every append would have to allocate a new array one slot longer and copy all the pointers into it, making a sequence of $n$ appends cost $O(n^2)$. Instead, when the array is full, CPython allocates a larger one with spare room, roughly one-eighth more than needed plus a little. The figure shows the real growth pattern of a CPython 3.12 list, measured with sys.getsizeof.

The capacity runs 0, 4, 8, 16, 24, 32, 40, 52, 64, 76, …: each reallocation adds about 12.5 % plus a constant. Most appends therefore find a free slot and cost almost nothing; occasionally one triggers a move. Press “append 10 times” repeatedly and watch the last readout: the average number of element copies per append stays bounded by a small constant as the list grows. Appending is therefore amortised $O(1)$: cheap on average, even though an individual append is occasionally expensive. The memory cost of the spare slots is modest, and growing by a fixed fraction rather than a fixed amount is what makes the total copying linear.

Lists are sorted in place with a.sort() or copied in sorted order with sorted(a), both using Timsort, an $O(n \log n)$ algorithm designed by Tim Peters for CPython that is especially fast on data that are already partly ordered. Both accept a key function, as in sorted(atoms, key=lambda a: a.mass). Checking membership with x in a, finding a position with a.index(x) and counting with a.count(x) scan the list and are $O(n)$.

Tuples and unpacking

A tuple is an immutable sequence, written with commas and usually parentheses: point = (0.0, 1.5, -2.0). Being immutable, a tuple can be a dictionary key or a set element, and it signals to a reader that its contents form a fixed record rather than a growing collection. Tuples are slightly smaller and faster to create than lists.

The commas, not the parentheses, make a tuple, and the reverse operation, unpacking, binds several names at once:

x, y, z = point                       # unpack a 3-tuple
a, b = b, a                           # swap two names, no temporary needed
first, *rest = [1, 2, 3, 4]           # first = 1, rest = [2, 3, 4]
for element, mass in [("H", 1.008), ("C", 12.011)]:
    ...
def stats(values):
    return min(values), max(values), sum(values) / len(values)
lo, hi, mean = stats(energies)        # several results from one function

A tuple whose fields deserve names can be made into a named tuple with collections.namedtuple, or, more flexibly, a data class (the chapter on classes), so that atom.mass replaces the cryptic atom[1].

Dictionaries

A dictionary maps keys to values: masses = {"H": 1.008, "C": 12.011, "O": 15.999}. Looking up masses["C"] takes the same short time whether the dictionary has three entries or three million. That is remarkable: with a list of pairs, finding a key requires scanning until it is found. The trick is a hash table.

Every hashable object has a hash, an integer computed from its value by the function hash(), with the property that equal objects have equal hashes. A dictionary keeps an array of slots whose size is a power of two. To store a key, it computes the key’s hash, takes the last few bits to choose a slot, and puts the key–value pair there. To find a key later, it repeats the calculation and goes straight to that slot. Different keys occasionally map to the same slot, a collision; the dictionary then probes further slots in a fixed sequence until it finds the key or an empty slot, which means the key is absent. The figure lets you build a table yourself.

Because collisions are rare when the table has plenty of empty slots, the dictionary keeps its load factor below two-thirds by doubling when it fills, just as a list over-allocates. The average lookup, insertion and deletion are therefore $O(1)$. The rules about keys follow directly from the mechanism. A key must be hashable, and its hash must never change while it is in the table, otherwise it would be looked for in the wrong slot; that is why keys must be immutable, so strings, numbers and tuples work and lists do not. And two keys that are equal must have equal hashes, which is why 1, 1.0 and True all count as the same key.

Since Python 3.7, dictionaries remember the order in which keys were inserted, and iterating over one yields keys in that order. The essential methods are d.get(key, default) (return a default instead of raising KeyError for a missing key), d.items(), d.keys() and d.values() for iteration, key in d for membership ($O(1)$), d.pop(key) and d.update(other). Two specialised dictionaries in collections remove boilerplate: defaultdict(list) creates an empty list automatically the first time a key is used, ideal for grouping, and Counter counts things.

Sets

A set is an unordered collection of distinct hashable objects, {"H", "C", "O"}, stored in a hash table without values. Its two uses are removing duplicates, set(symbols), and fast membership tests, x in s, which is $O(1)$ instead of the $O(n)$ of a list. Sets also support the operations of mathematics: union a | b, intersection a & b, difference a - b and symmetric difference a ^ b. Comparing the neighbour sets of an atom at two times, for example, shows which neighbours left and which arrived. An empty set is written set(), because {} is an empty dictionary.

The difference in membership cost is large enough to measure in the browser:

On a laptop running native Python the set is about ten thousand times faster for this test; in the browser the timer is coarser and the measured ratio is lower, but still in the thousands. The ratio grows in proportion to the size of the data.

Nested data and JSON

Real data are usually nested: a list of dictionaries, one per molecule, each holding a list of atoms, each atom a dictionary with an element and coordinates. Python’s containers nest freely, and access chains naturally: molecules[3]["atoms"][0]["xyz"][2]. Remember that nesting stores references: a dictionary placed into two lists is one object, and changing it changes both, as the chapter on objects and names showed. Use copy.deepcopy when independent copies are needed.

Nested dictionaries and lists of strings, numbers, Booleans and None map exactly onto JSON, the text format used by most web services and many materials databases. The json module converts in both directions:

import json
mol = {"name": "water", "atoms": [{"el": "O", "xyz": [0, 0, 0]},
                                  {"el": "H", "xyz": [0.96, 0, 0]}]}
text = json.dumps(mol, indent=2)       # to a string (or json.dump(mol, file))
back = json.loads(text)                # and back

For large numerical arrays, JSON is slow and bulky, and binary formats such as NumPy’s .npy, HDF5 or NetCDF are used instead.

Choosing a structure

The choice usually follows from the question the code will ask most often. If it asks “what is the i-th item?” or keeps items in order, use a list. If it asks “what is the value for this key?”, use a dictionary. If it asks “have I seen this before?” or “which items are in both?”, use a set. If the record is fixed and should not change, use a tuple or a data class. And if the items are numbers on which you will do arithmetic, use none of these: use a NumPy array, which stores raw numbers contiguously and operates on them in compiled code, as the next subject’s first topic explains.

References and further reading

Types of your own

Every object belongs to a type, and the built-in types, int, str, list, dict and the rest, cover a great deal. But the things a scientific program reasons about are atoms, molecules, crystal structures, potentials, trajectories and simulations. A class lets you define a new type whose objects bundle together the data describing such a thing (its attributes) and the operations that make sense on it (its methods). An Atom object might carry an element symbol, a mass and a position; a Potential object might know how to compute an energy and a force. Code that uses them can then speak the language of the problem, atom.mass and potential.force(r), rather than juggling parallel lists and remembering that index 3 means the mass.

This style, organising a program around objects that combine data and behaviour, is called object-oriented programming. Python supports it fully but does not force it: a short analysis script is often clearest as a handful of functions, and the decision to write a class should be made because it simplifies the code, not by habit. Classes earn their keep when several pieces of data always travel together and several functions always operate on them, when a family of related things should be interchangeable (several kinds of potential, several kinds of thermostat), or when you want your objects to work with Python’s own syntax, such as + for vectors.

Defining a class

A class is defined with the class statement:

class Particle:
    """A point particle in one dimension."""

    def __init__(self, mass, x, v=0.0):
        self.mass = mass
        self.x = x
        self.v = v

    def kinetic_energy(self):
        return 0.5 * self.mass * self.v ** 2

p = Particle(mass=2.0, x=0.0, v=3.0)
print(p.kinetic_energy())     # 9.0

Executing the class statement creates a class object, here bound to the name Particle, which holds the functions defined in its body. Calling the class, Particle(2.0, 0.0, 3.0), creates a new instance: an object of type Particle with its own empty set of attributes. Python then calls the special method __init__ on the new instance to initialise it, passing the instance as the first argument, conventionally named self, followed by the arguments of the call. The assignments self.mass = mass attach attributes to that particular instance. Every instance has its own attributes, so two particles can have different positions.

Methods are functions defined in the class body. When you call p.kinetic_energy(), Python finds the function kinetic_energy on the class and calls it with p as the first argument: p.kinetic_energy() is exactly Particle.kinetic_energy(p). That is all self is: the instance on which the method was called, passed explicitly. Unlike C++ or Java, Python requires you to write self.x inside methods to reach the instance’s attributes; a bare x would be a local variable.

The figure shows instances in memory. Each instance is a separate object with its own attributes, and the class itself is an object too, holding the methods.

How attribute lookup works

When you write obj.name, Python searches in a definite order. It looks first in the instance’s own attribute dictionary (obj.__dict__), then in the class, then in the class’s base classes. The first match wins. This rule explains several behaviours at once. Methods are found on the class because instances do not normally store them. Data stored on the class rather than on self, a class attribute, is shared by all instances: in the second example of the figure, every Atom sees unit = 'amu' through its class. Assigning o.unit = 'g/mol' does not change the class attribute; it creates a new attribute on the instance o that shadows the class one, so only o sees the new value.

The same rule produces a trap analogous to the mutable default argument of the chapter on functions. A mutable class attribute, such as atoms = [] written in the class body, is a single list shared by every instance; appending to it through one instance appends for all. Per-instance containers must be created in __init__, as the Molecule class in the third example does with self.atoms = [].

Special methods

Python’s syntax is defined in terms of methods with double-underscore names, called special or “dunder” methods. When the virtual machine executes BINARY_OP + on a + b, it calls a.__add__(b); len(x) calls x.__len__(); x[i] calls x.__getitem__(i); for item in x calls x.__iter__(); print(x) uses x.__str__() or x.__repr__(); a == b calls a.__eq__(b). By defining these methods, your own classes work with the language’s operators and built-in functions exactly like the built-in types. A small vector class shows the idea:

__repr__ deserves a special word: a class without one prints as something like <__main__.Vector object at 0x10a3f2e50>, which is useless in a debugging session. A good __repr__ returns a string that looks like the code that would recreate the object.

Special methods also explain why Python cares so little about the declared type of an object. A function that loops over its argument works with a list, a tuple, a file, a generator or any object of yours that defines __iter__. This is called duck typing: if it iterates like a sequence, it can be used as one. NumPy arrays, pandas data frames and PyTorch tensors all plug into Python’s syntax this way, which is why a + b adds two arrays element by element: the array type’s __add__ method runs a compiled loop.

Inheritance and polymorphism

A class can be defined as a specialisation of another, its base class, by naming the base in parentheses. The new class inherits every method of the base and can add methods or override existing ones. In simulation code this is a natural way to express a family of interchangeable models. Every pair potential provides an energy, and the force can be derived from any energy by differentiation:

The loop at the end treats both objects identically, calling energy and force without knowing which kind of potential it has: this is polymorphism, and it is what lets a simulation engine accept any potential, any thermostat or any integrator through a common interface. Harmonic and LennardJones define only energy; the finite-difference force is inherited from Potential (a real code would override it with the analytical derivative, which is faster and exact). Inside an overriding method, super().method(...) calls the base class’s version, which is the usual way for a subclass’s __init__ to run its parent’s initialisation before adding its own.

Inheritance is easy to overuse. Deep hierarchies of classes, where behaviour is spread across five levels of parents, are hard to follow. Often composition, where an object holds other objects as attributes, is clearer: a Simulation has a potential, an integrator and a thermostat, rather than being a special kind of any of them. The third example in the figure, a Molecule holding a list of Atom objects, is composition.

Data classes and properties

Many classes in scientific code are mostly records: an atom with an element, a mass and a position. Writing __init__, __repr__ and __eq__ by hand for each is tedious, and the dataclasses module generates them from annotated class attributes:

from dataclasses import dataclass, field

@dataclass
class Atom:
    element: str
    mass: float
    position: list = field(default_factory=lambda: [0.0, 0.0, 0.0])

o = Atom("O", 15.999)
print(o)                         # Atom(element='O', mass=15.999, position=[0.0, 0.0, 0.0])
print(o == Atom("O", 15.999))    # True: equality compares the fields

Note default_factory: even here, a mutable default must be created afresh for each instance. @dataclass(frozen=True) makes instances immutable and therefore usable as dictionary keys.

A property makes a method look like an attribute. It suits quantities computed from other attributes, so that they can never go stale:

class Box:
    def __init__(self, lx, ly, lz):
        self.lengths = (lx, ly, lz)

    @property
    def volume(self):
        lx, ly, lz = self.lengths
        return lx * ly * lz

box = Box(3.0, 3.0, 4.0)
box.volume          # 36.0, recomputed every time it is read

A small object-oriented simulation

The pieces come together in a class that holds the state of a system and advances it in time. The Oscillator below is a particle on a spring, integrated with the velocity Verlet algorithm that molecular dynamics uses for millions of atoms. The object stores the state, the force and energy methods describe the physics, and step advances the state by one time step.

After 10 000 steps, about 80 oscillation periods, the energy still agrees with its initial value to about one part in ten thousand, and the error oscillates rather than growing, a property of the Verlet algorithm explained in the chapter on integrators. Swapping the spring for another potential, or the integrator for another scheme, would mean changing one method, while all the code that runs and analyses the simulation stays the same. That separation of concerns is the real benefit of classes. Production codes such as ASE are organised exactly this way: an Atoms object holds positions, a Calculator object computes energies and forces, and a Dynamics object advances time.

References and further reading

Why NumPy exists

The Python data model is flexible, and the chapters on the language showed what that flexibility costs. A list of a million numbers is a million separate float objects scattered across memory, each 24 bytes with its own type pointer and reference count, plus a million 8-byte pointers in the list. Adding two such lists element by element runs a million iterations of the interpreter loop, each checking types, dispatching an addition and allocating a new object. For numerical work, that is hopelessly slow and wasteful.

NumPy solves the problem with a single new type, the n-dimensional array or ndarray. An array stores numbers the way Fortran and C do: as raw binary values of one fixed type, packed side by side in one contiguous block of memory. Operations on arrays are implemented in compiled C loops that run over this block without ever creating Python objects for the individual numbers. A Python program that uses NumPy well pays the interpreter’s overhead once per operation rather than once per element. Almost the entire scientific Python ecosystem, SciPy, pandas, Matplotlib, scikit-learn, ASE and MDAnalysis among many others, is built on NumPy arrays, and the array interface of PyTorch and JAX was designed to imitate it.

The figure at the end of this chapter measures the difference in your browser. Before that, this chapter explains how an array is laid out in memory, because nearly every rule of NumPy, from which operations copy data to why some loops are faster than others, follows from it.

An array is a buffer with a shape, a dtype and strides

An ndarray consists of two parts. The data buffer is a block of raw bytes. The header describes how to interpret those bytes as an array, with four pieces of information:

For a (3, 4) array of float64 stored in the default C order (row-major), the twelve values are laid out row after row. Moving one column to the right moves 8 bytes; moving one row down skips a whole row of 4 elements, 32 bytes. The strides are therefore (32, 8). In Fortran order (column-major), used by Fortran and by many linear-algebra libraries, the values are laid out column after column, and the strides are (8, 24).

Strides make it possible to represent many different arrays over the same buffer without copying anything. Choose a[:, 2] in the figure: the column is described by the same buffer, a starting offset of 16 bytes and a single stride of 32 bytes, jumping from row to row. Choose a.T: the transpose has the shape reversed and the strides swapped, (8, 32), and not a single byte has moved. a[::2, ::2] doubles both strides. a.reshape(4, 3) on a C-ordered array just reinterprets the same bytes with a new shape. Such arrays are called views. They are created almost instantly, regardless of size, and they share memory with the original: writing to a view writes to the original array.

Not every selection can be expressed with strides. Choosing the elements where a condition holds, a[a % 3 == 0] in the figure, picks an irregular set of positions, so NumPy must copy the selected values into a new buffer. The same is true of indexing with a list of positions, a[[0, 2]], called fancy indexing. Knowing which operations give views and which give copies avoids two kinds of bug: modifying an array through a view without meaning to, and modifying a copy and expecting the original to change. When in doubt, np.shares_memory(x, y) answers the question, and .copy() forces an independent array.

Memory layout also affects speed. Processors read memory in cache lines of 64 bytes, so a loop that walks through consecutive addresses uses every byte it fetches, while a loop that jumps by large strides wastes most of each fetch (the chapter on computer architecture shows the memory hierarchy). Summing along the rows of a large C-ordered array is therefore faster than summing along its columns, and copying a transposed array before a long computation can pay for itself.

Vectorised arithmetic

Arithmetic operators on arrays act element by element: a + b, a * b, a ** 2, np.sqrt(a) and np.exp(-a / kT) each produce a new array of results, computed by a compiled loop over the buffers. These element-wise functions are called universal functions or ufuncs. Writing a computation as whole-array operations instead of a Python loop is called vectorisation, and it is the central skill of numerical Python.

x = np.linspace(0.0, 2.0 * np.pi, 1000)   # 1000 evenly spaced points
y = np.sin(x) * np.exp(-x / 3.0)          # 1000 values, no Python loop

Reductions combine many elements into fewer: a.sum(), a.mean(), a.std(), a.min(), a.max(), a.argmax(). With the axis argument they reduce along one axis only. For a trajectory stored as an array of shape (n_frames, n_atoms, 3), positions.mean(axis=0) gives the average position of each atom over time, with shape (n_atoms, 3), while positions.mean(axis=1) gives the centre of each frame, with shape (n_frames, 3). A useful way to remember it: the axis you name is the one that disappears.

Linear algebra has its own operators. a @ b is matrix multiplication, which for large matrices calls an optimised BLAS library that uses vector instructions and multiple cores and can approach the processor’s peak speed. np.linalg provides solve, inv, eigh (eigenvalues of symmetric matrices, as in the chapter on matrices), svd and norm.

Indexing and slicing

Arrays are indexed with one index per axis, separated by commas: a[1, 2] is the element in row 1, column 2. Slices work on every axis, a[:, 1:3] selecting all rows and columns 1 and 2, and an omitted trailing index means “everything”, so a[1] is row 1. Negative indices count from the end. As explained above, basic slicing always returns a view.

Two kinds of advanced indexing return copies. Boolean indexing uses an array of True/False of the same shape as a mask: energies[energies < cutoff] selects the low energies, and r[r < r_cut] selects the distances inside a cut-off. Combined with assignment, a mask modifies selected elements in place: v[np.abs(v) > vmax] = 0.0. Fancy indexing uses arrays of integer positions: positions[[0, 5, 9]] picks three atoms, and np.argsort(e) returns the positions that would sort an array, so structures[np.argsort(e)[:10]] gives the ten lowest-energy structures.

Broadcasting

Arithmetic between arrays of different shapes is allowed when the shapes are compatible according to the broadcasting rules, which let a smaller array act as if it were repeated to match a larger one without actually copying it.

The stretching is done with a stride of zero: a size-1 axis repeated $n$ times is described by stepping zero bytes along it, so the same value is read again and again. No memory is used for the virtual copies. The figure below applies the rules to any pair of shapes and draws the result for two-dimensional cases.

Broadcasting turns many loops into single expressions. Subtracting the centre of a molecule from every atom is positions - positions.mean(axis=0): shapes (N, 3) and (3,) broadcast to (N, 3). An outer sum, x[:, None] + y[None, :], uses None (also spelt np.newaxis) to insert size-1 axes, turning shapes (n,) and (m,) into (n, 1) and (1, m) and producing an (n, m) table. The same trick computes all pair vectors between atoms at once, which is the core of the example at the end of this topic.

Fixed-width numbers

NumPy’s numbers behave like those of C and Fortran, not like Python’s. Integers have a fixed width and overflow silently: an int8 holds values from −128 to 127, so multiplying 2 by 100 in int8 gives −56, and even int64 arithmetic wraps around above $2^{63} - 1$. Python’s own integers never overflow, so this surprises people moving data between the two. Choose dtypes deliberately: float64 for physical quantities unless memory or GPU speed demands float32, and integer types wide enough for the counts involved.

Floating-point types carry the rounding behaviour of IEEE 754. float32 keeps about seven significant digits, which is enough to store coordinates but not to accumulate a sum of a million energies without visible error. NumPy’s sum reduces this risk by adding in pairs (pairwise summation), but long accumulations, such as running averages over a simulation, should be done in float64. Division by zero produces inf or nan with a warning rather than an exception, and nan then propagates through every later calculation; np.isfinite(a).all() is a cheap sanity check after a long computation.

How much faster, and why

The figure below times four ways of computing $\sum_i x_i^2$ for up to a million numbers, in this browser: a Python for loop over a list, the built-in sum of a generator, the NumPy expression (a * a).sum(), and the dot product a @ a, which hands the whole job to a single BLAS routine.

On a laptop running native Python, the NumPy versions of this simple sum are about 15 to 100 times faster than the loop, and the advantage is larger for heavier arithmetic such as exponentials and square roots. In the browser, where NumPy runs as WebAssembly with fewer vector instructions, the gap is smaller, typically ten- to twenty-fold. Three effects add up. NumPy executes one interpreter-level operation instead of millions, so the per-bytecode overhead disappears. Its compiled loop reads raw doubles from consecutive addresses, which the processor’s caches and prefetcher handle at full speed. And the compiled loops use SIMD (single instruction, multiple data) instructions that process two, four or eight numbers per machine instruction. The last line, the dot product, also avoids creating the temporary array a * a, which is why it is faster still.

Vectorisation has two costs to keep in mind. Temporary arrays use memory: np.exp(-(r / sigma)**2) creates three intermediate arrays the size of r. For very large arrays this can exhaust memory, and in-place operations (r /= sigma, np.exp(r, out=r)) or chunking help. And some algorithms are inherently sequential, such as a time-stepping loop where each step needs the previous one; for those, the loop over time stays in Python (or moves to compiled code, as in the chapter on performance), while the work over atoms inside each step is vectorised.

A worked example: all pair distances

Computing the distances between all pairs of $N$ atoms is the first step of almost every structural analysis, from coordination numbers to radial distribution functions. With loops it is a double loop over pairs. With broadcasting, inserting a new axis in two different places turns the (N, 3) position array into shapes (N, 1, 3) and (1, N, 3); their difference broadcasts to (N, N, 3), every pair vector at once.

The line d -= L * np.round(d / L) applies the minimum-image convention of periodic boundary conditions, explained in the chapter on molecular simulation, to all 250 000 pair vectors in one expression. The approach uses $O(N^2)$ memory, which is fine for a few thousand atoms; beyond that, the pairs are found with neighbour lists or with scipy.spatial.cKDTree, which only examine nearby atoms.

References and further reading

From raw output to a result

A typical computational study produces files: trajectories, energy logs, structure files, tables of computed properties. Turning those files into a number with an error bar, or a figure in a paper, follows the same sequence almost every time. The files are read and parsed into arrays; the arrays are cleaned and checked; a quantity is computed or fitted; its uncertainty is estimated; and the result is plotted and written out. Each step is easy on its own. Doing all of them correctly, in a way that can be rerun next month on new data and give the same answer for the same input, is the real skill, and Python’s scientific libraries make it straightforward. This chapter goes through the steps in order.

Files and paths

A file is opened with the built-in open, which returns a file object. The safest way to use it is inside a with block, which closes the file automatically when the block ends, even if an error occurs inside it:

with open("energy.dat") as f:            # mode "r" (read, text) is the default
    for line in f:                       # a file is an iterator over its lines
        if line.startswith("#"):
            continue
        t, e = map(float, line.split())

The second argument of open selects the mode: "r" to read, "w" to write (erasing any existing file), "a" to append, and a "b" suffix for binary files, which return bytes instead of text. Text files are decoded with an encoding, almost always UTF-8; specifying encoding="utf-8" avoids surprises on systems with other defaults. Iterating over the file object reads one line at a time, so even a file of many gigabytes can be processed with constant memory, as the trajectory generator of the chapter on control flow did.

Paths are best handled with pathlib.Path, which joins directory names portably, finds files by pattern and reads small files in one call:

from pathlib import Path
runs = Path("simulations")
for log in sorted(runs.glob("T*/md.log")):      # every md.log in T300/, T310/, ...
    temperature = int(log.parent.name[1:])
    text = log.read_text()

Many structure formats are simple enough to parse by hand. An XYZ file contains the number of atoms, a comment line, and one line per atom with its element and coordinates:

For anything more complicated, use a library that already handles the format’s corner cases: ASE reads and writes dozens of structure formats (ase.io.read("POSCAR")), and MDAnalysis and MDTraj read trajectories from all the common simulation engines.

Tables of numbers

Most numerical output is a table: columns of numbers separated by spaces or commas, with some header lines. np.loadtxt reads such a file directly into an array. GROMACS writes its .xvg files with header lines beginning with # and @; telling loadtxt to treat both as comments is all that is needed:

np.genfromtxt does the same but tolerates missing values, and np.savetxt writes an array back out with a header. For arrays that only Python will read, the binary formats are far better: np.save("traj.npy", positions) writes the exact bytes of the buffer with a small header recording the dtype and shape, so loading is instant and lossless, and np.savez_compressed stores several named arrays in one file. Large simulation data sets are usually kept in HDF5 (via h5py) or NetCDF, which add compression, partial reading and metadata.

When a table has named columns of mixed types, such as a list of materials with their structure type and computed properties, pandas is the right tool. Its DataFrame is a table whose columns are NumPy arrays, with labels and a large set of operations for selecting, grouping, joining and summarising:

pd.read_csv reads comma- or space-separated files with headers, guessing column types, and df.to_csv writes them. A good rule is to use NumPy for homogeneous numerical arrays, such as coordinates, trajectories and grids, and pandas for heterogeneous tables of records, such as the results of a high-throughput screening.

Fitting models to data

Fitting determines the parameters of a model that best describe some data, and it is how most physical constants are extracted from simulations. The commonest method is least squares: choose parameters that minimise the sum of squared differences between the data and the model’s predictions. The chapter on probability explains why this is the right choice when errors are independent and Gaussian, and learning from data develops it as the first machine-learning method.

For a model that is linear in its parameters, such as a straight line or a polynomial, least squares has an exact solution, and np.polyfit computes it together with the covariance matrix of the parameters, whose diagonal gives their variances. The Arrhenius law, $k = A e^{-E_a/RT}$, becomes a straight line when plotted as $\ln k$ against $1/T$, with slope $-E_a/R$:

For a model that is nonlinear in its parameters, such as an exponential decay with an offset, scipy.optimize.curve_fit minimises the squared residuals iteratively from an initial guess and again returns parameters and their covariance:

from scipy.optimize import curve_fit

def model(t, a, tau, c):
    return a * np.exp(-t / tau) + c

params, cov = curve_fit(model, t, y, p0=[1.0, 1.0, 0.0])
errors = np.sqrt(np.diag(cov))

Three cautions apply to every fit. The reported uncertainties assume that the residuals are independent and that the model is correct; for correlated simulation data, such as consecutive points of a time series, they are too small, and block averaging or bootstrapping gives honest errors (trajectory analysis). Nonlinear fits can converge to a wrong local minimum if the initial guess is poor, so start from physically sensible values. And always plot the residuals: a systematic pattern in the differences between data and fit is the clearest sign that the model is wrong, which no error bar will reveal.

Plotting with Matplotlib

Matplotlib is the standard plotting library. Its object model has two main pieces: a Figure, the whole image, and one or more Axes, each a single plot with its own coordinate system. The recommended interface creates both explicitly and then calls methods on the axes:

import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(4, 3))        # size in inches
ax.plot(x, y, "o", label="simulation")
ax.plot(x_fit, y_fit, "-", label="fit")
ax.set_xlabel("time (ps)")
ax.set_ylabel("energy (kJ/mol)")
ax.legend()
fig.savefig("energy.pdf", bbox_inches="tight")

The figure below runs real Matplotlib in this page and shows the image it produces. Edit the code and run it again: change the colours, add a second panel with plt.subplots(1, 2), or switch to logarithmic axes with ax.set_yscale("log").

A good scientific figure has labelled axes with units, symbols large enough to read when the figure is shrunk to one column of a journal page, a colour scheme that survives printing in greyscale and colour-blind readers (Matplotlib’s default palette and the viridis colour map are designed for this), and no decoration that carries no information. Save figures for papers in a vector format, PDF or SVG, so that lines stay sharp at any size, and raster images (PNG) at 300 dpi or more for slides and web pages.

A complete small analysis

The steps combine into a complete analysis of a simulated system. Brownian particles take independent Gaussian steps; the mean-squared displacement of a diffusing particle grows linearly in time, $\langle |\mathbf r(t) - \mathbf r(0)|^2\rangle = 6Dt$ in three dimensions, so a straight-line fit to it gives the diffusion coefficient $D$. Generating the trajectories, computing the MSD and fitting take a few lines of NumPy, and a known input value lets us check the whole pipeline:

The estimate comes out within a couple of per cent of the true value of 1. Testing an analysis on synthetic data with a known answer in this way, before applying it to real simulation output, is the most effective guard against subtle mistakes in units, indexing or averaging.

Organising an analysis

A few habits turn a collection of scripts into a reliable analysis. Keep each project in its own environment with the package versions recorded, so that it can be recreated. Explore in a Jupyter notebook if you like, but move code that works into functions in .py modules and import them, so that the same function is used everywhere and can be tested. Keep raw data read-only and separate from derived results, and write every derived file from a script rather than by hand, so that the whole chain from raw output to figure can be rerun with one command. Seed random-number generators, as the examples here do with default_rng(42), so that stochastic results can be reproduced exactly. And put the code under version control with git, so that every figure can be traced back to the exact code that made it.

References and further reading

Measure before optimising

Most code is fast enough, and most of the time spent in slow code is spent in a small part of it. Donald Knuth’s warning that premature optimisation is the root of much evil is about exactly this: effort spent speeding up code that does not dominate the run time is wasted, and usually makes the code harder to read. The first step is therefore always to measure, both to find where the time goes and to know whether a change helped.

For timing a small piece of code, the timeit module runs it many times and reports the best time, which removes noise from other activity on the machine. In Jupyter the magic %timeit expression does the same. For finding where time goes in a whole program, a profiler records how often each function is called and how long it takes. The standard library’s cProfile does this with little overhead:

The report shows the number of calls to each function and the time spent inside it (tottime) and inside it plus everything it calls (cumtime). Here the 76 636 calls to distance, with its generator expression, dominate. For a line-by-line view, the third-party line_profiler package annotates every line of chosen functions with its time. Whatever tool is used, the result usually points to one loop, and that loop is where the techniques below should be applied.

Better algorithms first

The largest speed-ups come from doing less work, not from doing the same work faster. Reducing an $O(n^2)$ algorithm to $O(n \log n)$ or $O(n)$ beats any constant-factor improvement once $n$ is large (the chapter on collections introduced this notation). In simulation code the classic example is finding neighbours: checking every pair of $N$ atoms costs $O(N^2)$, but dividing the box into cells no smaller than the cut-off and only comparing atoms in neighbouring cells costs $O(N)$, a factor of a thousand for a million atoms. In Python, scipy.spatial.cKDTree provides fast neighbour searches of this kind.

Smaller algorithmic changes matter too. Membership tests in a loop should use a set or dictionary, not a list. Repeated work can be cached: if an expensive function is called many times with the same arguments, decorating it with @functools.lru_cache stores previous results and returns them instantly. Quantities that do not change inside a loop should be computed once outside it. Sorted data can be searched in $O(\log n)$ with bisect or np.searchsorted instead of scanning.

Vectorise

The next step is to move the loops out of the interpreter. As the chapter on NumPy arrays explained and measured, a vectorised NumPy expression runs a compiled, cache-friendly loop and is commonly tens to hundreds of times faster than the equivalent Python loop. The energy calculation above vectorises in a few lines:

The result agrees with the pure-Python version to twelve significant digits, and the time falls by about a factor of ten, from roughly 80 ms to under 10 ms on a laptop; the gap widens as the number of atoms grows. Vectorisation is limited by memory, since it creates temporary arrays, and by algorithms with sequential dependencies, such as time stepping, where each step needs the result of the previous one. For those cases the remaining loop must be compiled.

Compile the hot loop

Several tools turn a Python-like loop into machine code, removing the interpreter’s per-instruction cost entirely.

Numba is the least effort. Decorating a function that uses NumPy arrays and simple loops with @numba.njit makes Numba compile it, the first time it is called, into machine code specialised for the types of the arguments, using the LLVM compiler infrastructure. Inside, loops over array elements run at speeds close to C or Fortran:

import numba
import numpy as np

@numba.njit(cache=True)
def total_energy(pos):
    n = pos.shape[0]
    e = 0.0
    for i in range(n):
        for j in range(i + 1, n):
            r2 = 0.0
            for k in range(3):
                d = pos[i, k] - pos[j, k]
                r2 += d * d
            sr6 = 1.0 / (r2 * r2 * r2)
            e += 4.0 * (sr6 * sr6 - sr6)
    return e

The explicit triple loop, which would be very slow in plain Python, is now fast, and it uses no temporary arrays at all. Numba also offers parallel=True with numba.prange to spread a loop over cores, and it can compile for GPUs. Its limitation is that it supports only a subset of Python: numbers, NumPy arrays and simple data structures, not arbitrary objects or most libraries.

Cython compiles a superset of Python with optional C type declarations into a C extension module; it is how much of SciPy and scikit-learn is written. f2py, part of NumPy, wraps Fortran subroutines so that they can be called from Python with NumPy arrays as arguments, which makes it easy to reuse existing Fortran kernels or to write new ones (the chapter on Fortran shows how). For maximum control, extensions can be written in C, C++ (with pybind11) or Rust (with PyO3). All of these follow the same design principle: keep the high-level logic, input and output in Python, and put the few loops that dominate the run time in compiled code.

The interpreter itself is getting faster, too. CPython 3.11 introduced the specialising adaptive interpreter described in the chapter on how Python runs, which made typical code around 25 % faster, and CPython 3.13 added an experimental just-in-time compiler. Alternative interpreters such as PyPy can run unmodified pure-Python loops several times faster, although they are less compatible with compiled extensions like NumPy.

Memory

Speed and memory are linked, because moving data is often slower than computing with it. Choosing the smallest adequate dtype halves memory traffic when moving from float64 to float32, which is why GPUs and machine-learning models often use single or even half precision, at the cost of the rounding issues discussed in the chapter on numbers. In-place operations (a *= 2, np.add(a, b, out=a)) avoid allocating temporaries. Generators and chunked processing keep the working set small: a trajectory too large for memory can be processed one frame at a time, or opened as a memory-mapped array with np.load("traj.npy", mmap_mode="r"), which reads only the parts that are actually touched. Access patterns matter as well: walking through an array in its memory order uses the cache efficiently, while jumping across it with large strides does not.

Parallelism

A modern workstation has 8 to 64 cores and a cluster node often more, so running work in parallel can multiply throughput. How much it helps is limited by the part of the work that cannot be parallelised, a fact captured by Amdahl’s law:

With 5 % serial work, no number of cores can give more than a 20-fold speed-up, and 128 cores give only about 17. Communication between cores adds a cost that grows with their number, so beyond some point adding cores makes a job slower. The figure lets you explore both effects.

Python adds one complication of its own. In standard CPython, the global interpreter lock (GIL) allows only one thread at a time to execute Python bytecode. Threads therefore do not speed up pure-Python computation, although they are useful for waiting on input and output, and they do help when the work happens inside compiled code that releases the lock, as NumPy’s large array operations and BLAS routines do. CPython 3.13 introduced an optional free-threaded build without the GIL (PEP 703), which is gradually becoming practical as libraries adapt to it.

For CPU-bound Python code, the standard approach is to use processes instead of threads. Each process has its own interpreter and memory, so there is no lock to share. concurrent.futures.ProcessPoolExecutor, multiprocessing.Pool and the joblib library make it easy to apply a function to many independent inputs in parallel:

from concurrent.futures import ProcessPoolExecutor

def analyse(path):
    ...                                   # read one trajectory, return a number
    return result

if __name__ == "__main__":
    with ProcessPoolExecutor(max_workers=8) as pool:
        results = list(pool.map(analyse, trajectory_files))

Many scientific workloads are embarrassingly parallel in exactly this way: independent simulations at different temperatures, analysis of separate trajectories, single-point calculations on thousands of structures. The best way to run them on a cluster is often not inside Python at all, but as many separate jobs submitted to the scheduler, for example as a Slurm job array (the chapter on Linux for simulation). For tightly coupled parallel programs that span many nodes, the standard is MPI, available in Python through mpi4py, where each process works on part of the problem and they exchange data by passing messages.

Two practical points prevent common surprises. NumPy’s linear algebra already uses several threads through its BLAS library; running eight such processes on an eight-core machine can oversubscribe it with 64 threads, which is slower than either alone. Setting OMP_NUM_THREADS=1 (and the equivalents for OpenBLAS or MKL) in each worker avoids this. And GPUs offer another level of parallelism for array arithmetic: CuPy provides a NumPy-compatible array library on NVIDIA GPUs, and PyTorch and JAX run array computations and automatic differentiation on GPUs, which is how machine-learning potentials and differentiable simulations reach their speed.

A checklist

When a Python program is too slow, the following order of steps usually gives the most improvement for the least effort. Measure with a profiler to find the few functions that dominate. Look for a better algorithm or data structure for them: neighbour lists instead of all pairs, sets instead of lists, caching instead of recomputation. Vectorise the inner loops with NumPy. If a loop cannot be vectorised, compile it with Numba, or with Cython or f2py when Numba’s subset is too restrictive. Reduce memory traffic with appropriate dtypes and in-place operations. Then parallelise: independent tasks across processes or cluster jobs, array arithmetic on GPUs, and tightly coupled computations with MPI, checking scaling with a small test before committing many cores. At each step, keep the original version and compare results, because a fast program that gives wrong answers is worse than a slow one.

References and further reading