My day job, such as it is, involves optimizing simulations and scientific tooling. Right now I'm working on an astrophysical simulation with a group at CUNY. One of the grad students has spent the last few months working on a tool to process the output of a simulation.

The output is pretty big, about 200GB for a local test run, probably going into the tens of terabytes once we do a big run on the actual cluster. The tool, which needs to walk the output tables to construct a tree of black hole mergers (and other events, but that's the focus), takes about an hour to run. I answered a request to take a look and see what can be done to speed it up.

The first thing I discover is that the data is split across tens of thousands of .txt files, with every simulation timestep producing several individual tables, some of them just a few rows, others tens of thousands of rows.

The second thing I notice is that the postprocessing code is looping over the files, reading the same one multiple times in some places, to extract the merger information. There's a gigantic dictionary of dictionaries, and the leaf dictionaries are manually simulating a binary tree using labels as keys ("root", "A", "B", "A1", "A2", etc), and the values are pandas dataframes.

It is very difficult to understand. I am reminded that astrophysics grad students have a great deal of brainpower and focus available to them, and this may sometimes be counterproductive.

There's no good way to optimize this code in-place. Luckily, I've built some credibility with this team, so they trust me when I basically tell them we're going to redo the tree construction code entirely, and try to preserve the graphing code, which is also elaborate but is a lot saner, mostly just trying to work within networkgraph's preferred input mode.

Still, I get them to turn their runner into a script rather than a jupyter notebook, install snakeviz, and get a profile on a short run. They audibly gasped when they saw the snakeviz window pop open, and they could just see where all the time was going (mostly walking lots of small in-memory dataframes and loading them from txt). Thing is, they've seen this tool before, I've showed it in previous meetings for other performance work on the main simulation, they just thought it was something I coded up manually rather than something they can also do with cPython's bundled profiler and a single pip install.

It really shouldn't be shocking, since about 2 years ago I was in pretty much the same place, but it surprised me anyway. I sometimes think we need a 'Missing Semester of Your CS Education' equivalent specifically for scientists who got introduced to Python + data science tools and use that for pretty much everything. Intro to profilers and debuggers, the python memory model, useful/harmful data structures, and when NOT to use a dataframe. I think it would be useful.