Numerical Python

1 Arrays in Python

For many developers, numpy is the default number wrangler. It is easier than FORTRAN but not actually that easy, because its gradually-evolved design is not that great. In practice,lots of development time goes not on the problem itself, but on debugging what the hell jsut went wrong with the ant-nest of punctuation that you jsut dumped on the page that was supposed to do something fancy to an array.

Cas in point: NumPy’s array broadcasting. Broadcasting describes how numpy treats arrays with different shapes during arithmetic operations, allowing for efficient, C-speed vectorized operations without making explicit copies of data. Sometimes it works and feels like magic. Other times it works, but at the price of a baffling bunch of invocations of confusingly named helper functions. When it fails, it leads to confusing errors or, even worse, incorrect results that fail silently.

My favourite introduction to this is the DYNOMIGHT post “DumPy: NumPy except it’s OK if you’re dum”. The article argues that numpy’s syntax can be powerful, but irritating leading to code that is hard to read and debug. For example, to add a 1-D array to a 2-D array, a developer might need to insert a new axis with np.newaxis or None, like matrix + vector[:, None]. Is this the syntax we woudl have designed for ourselves? The author’s dumpy library proposes solutions that are more explicit, offering functions like dumpy.add_vec_to_cols(matrix, vector) or dumpy.add_vec_to_rows(matrix, vector). This approach makes the code’s intent unambiguous and self-documenting, getting the grug-brained seal of approval. The fact that dumpy’s implementation took only 700 lines of code is some kind of vindication of python, I guess.

I’ll confess that I don’t actually use dumpy myself; instead I try to make most of my numeric computations happen using einstein operations.

2 Pretty printing numbers

Getting the numbers to show up on the screen in a way that doesn’t make my eyes bleed is a constant, low-level chore. The defaults are rarely what I want. This annoys everyone and so there are many confusingly documented tools in the standard library and then many other libraries besides.

# Format x with 4 decimal points, aligned to fill 10 columns
print(f"{x:10.4f}")

import numpy as np
np.set_printoptions(edgeitems=5,
  linewidth=85, precision=4,
  suppress=True, threshold=500)

3 Einstein operations

Lots of array operations end up being easily expressed by shorthand from physics: the Einstein summation convention, or einsum for short if you are writing code. The core idea is simple: if an index variable appears twice in a term, it implies summation over all possible values of that index. So, instead of writing the dot product of two vectors a and b with an explicit sum symbol, like $\sum _i{a}_i{b}_i$, you can just write $a_i{b}_i$. The repeated index i tells you everything you need to know.

This convention elegantly generalizes to more complex operations like matrix multiplication, transpositions, and tensor contractions. It’s a compact and powerful way to describe operations on multi-dimensional arrays. For a deeper dive, there are some excellent tutorials out there:

• Einsum is All You Need by Tim Rocktäschel.
• einsum - an underestimated function by Chris Lemke.
• Understanding einsum for Deep learning which uses it to build a transformer from scratch.
• the einops manual illustrates everything with pictures which was what made the whole click for me.

np.einsum(’ik, kj->ij’, A, B)

rearrange(images, 'b h w c -> b c h w')

4 Optimisation

See optimisation.

5 Compiling

Think you can go faster by compiling some numeric code and invoking it from python? Or compiling some python subset?

Specific CPU and GPU support in Python precompiled binaries is weak and outdated. For maximum performance we should be compiling them.

e.g. for performance or invoking external binaries.

See compiling Python.

6 HDF5: h5py and PyTables

I use h5py frequently for saving and loading large arrays. It’s a straightforward, Pythonic interface to the HDF5 file format, and it feels like a natural extension of NumPy. My data gets saved in a hierarchical, self-describing format, which is great. The only quirk I’ve run into is some occasional weirdness when passing serialized h5py objects on macOS, which doesn’t happen on Linux.

I’ve also bookmarked PyTables, another HDF5 interface, but I’ve never been entirely clear on its value proposition over h5py. After looking into it, the distinction seems to be that h5py is a direct, NumPy-centric mapping to the HDF5 library, whereas PyTables builds more on top of it. PyTables adds a database-like layer, offering features like indexing and querying for searching within datasets. For my typical use case of just saving and loading entire arrays, h5py feels simpler. But if I had a massive HDF5 file and needed to efficiently query for specific rows or metadata without loading the whole thing, PyTables sounds like itmight add some value?

Much more info about this at data file formats in general, plus also see my hdf-specific notebook.

7 zarr

Zarr is a format for storing chunked, N-dimensional arrays that I’ve bookmarked because it feels like a modern take on HDF5. Its main value proposition seems to be its flexibility and cloud-friendliness. Unlike HDF5, which is a more monolithic format, zarr can store array chunks in lots of different backends—on disk, in a Zip file, or, most interestingly, directly in a cloud object store like Amazon S3.

This design makes it particularly good for parallel and distributed computing. It’s built to allow multiple threads or processes to read and write to the same array concurrently, which is a known pain point with some HDF5 configurations. While HDF5 has broader support in other languages like C++, zarr seems like the more forward-looking choice for pure Python, cloud-based workflows.

8 Dask

Dask is a library I’m aware of for parallel computing in Python. It’s interesting because it tackles the “big data” problem by mimicking the APIs of libraries I already use every day. It provides parallel versions of NumPy arrays and Pandas DataFrames. This means I can, in theory, apply the same code I wrote for in-memory data to datasets that are too large to fit on my machine, and Dask handles the distribution and parallel execution under the hood.

It’s built on two core components: a dynamic task scheduler for coordinating the work, and the “big data” collections themselves. This makes it flexible enough for simple, single-machine parallelism (taking advantage of all my CPU cores) all the way up to large, multi-machine clusters. For some examples of what this looks like in practice, Matthew Rocklin has a helpful how-to guide.

9 xarray

xarray solves a problem that has bitten me many times: keeping track of what the dimensions of my NumPy arrays actually mean. xarray wraps raw NumPy arrays and lets me attach labels to the dimensions (like time, latitude, longitude) and coordinates to the axes, which is useful for statistical analysis.

This makes the code more intuitive, self-documenting, and less prone to bugs. Instead of data.sum (axis=2), I can write data.sum (dim=‘channel’) This labeled approach is borrowed from the netCDF file format, which is a standard in the geosciences. I’ve also seen it popping up more in deep learning workflows, especially for handling complex, multi-dimensional data like satellite imagery or climate simulations, where keeping track of axes by number alone is a recipe for disaster. It seems to be a key tool for bridging the gap between raw numerical data and machine learning frameworks.

10 References