[jmschrei]

This optimization was found with auto-optimize, a tool developed by the Programmable Genomics Laboratory at UMass Chan Medical School: it speeds up bonsai_main.py --step all by 4.2x to 7.9x on every input timed except Bonsai's 64-cell example (3.2x and 3.6x, on runs of about 3 s), with output byte-identical to main. See How the output was checked.

On the 500-cell PBMC input at 4 MPI processes, the median wall time falls from 80.22 s to 13.41 s (5.98x, six paired runs from a fresh clone); on the four full 10x Genomics experiments the geometric mean speedup is 5.99x, and pbmc_10k_v3 (11,769 cells) at 16 processes takes 525.62 s instead of 3,526.12 s. Most of the gain comes from three things: numba kernels for the per-gene likelihood loops, freezing the garbage collector after the imports, and calling scipy's L-BFGS-B routine without minimize's per-call set-up. Peak memory of the largest process is 0.79 to 1.28 times main's (median 1.02; see Memory), and no OS-specific calls are added.

It builds on #99, which adds a pytest suite and changes no source. Of this branch's seven commits, the first six are #99's; this pull request's own change is the last one: 8 files, +2,368 / -186. Once #99 is merged, only that commit remains.

Changes

Roughly in the order the pipeline reaches them:

1. Garbage collector frozen after the imports (bonsai/bonsai_main.py): one gc.freeze(), so Bonsai's explicit gc.collect() calls no longer rescan every object the imported libraries leave alive.
2. MpiInfo built once (bonsai/mpi_wrapper.py): the namedtuple class is defined at import instead of on every get_mpi_info() call. scatter() is added beside bcast() and gather().
3. Star-tree times on every process (optimiseTStar in bonsai/bonsai_treeHelpers.py, SCData.__init__ in bonsai/bonsai_dataprocessing.py): each step of the EM-like scheme computes the children's times on all MPI processes instead of process 0 alone, in preprocessing and in the periodic re-optimisation during merging.
4. Three-leaf star-tree objectives (getLogLikAndGradStarTreeSequentialWrapper, getLogLikAndGradStarTreeWrapper): the objective and gradient that L-BFGS-B evaluates for every candidate pair are one numba kernel.
5. L-BFGS-B without minimize's set-up (new bonsai/bonsai_lbfgsb.py): the three optimisation sites call scipy's setulb directly, with the arguments and evaluation points minimize uses. It runs only on scipy 1.17.1, the version requirements.txt pins; elsewhere minimize runs as before.
6. A pair's star tree and closing log-likelihood (calcSingleDLogL): two kernels.
7. Upper-bound derivative terms (estimateDerBasedDLogLUB): one kernel; the norms stay with numpy.
8. Waiting processes evaluate merge pairs (MergePairPool in bonsai_treeHelpers.py): in the merge rounds process 0 used to run alone, processes 1 and up compute batches of candidate pairs. Every decision stays on process 0, in task order.
9. Known upper bounds looked up in a dict (getNewPairs): membership is tested in a dict instead of a list.
10. Pairs ordered by upper bound read on demand (PairSequence): getNewPairs no longer converts the whole array of pairs into a list of tuples at four places.
11. Nearest-neighbour connections as sets (NNConnectivity): replaces the scipy.sparse.lil_array, with the same entries, row reads and errors.
12. A node's position and log-likelihood from its children (findNodeLtqsGivenLeafs, getLtqsComplete): one kernel for nodes with 2 to 4 children.
13. Tree.optTimes' objective on a table (_OptTimesObjective): the tree is read into arrays once per call, and each evaluation is two kernels instead of a recursion over the nodes.
14. Two-leaf branch-length search (getOptTime2LeafTree, get_dlogl_spr_move): the bracket search and a port of scipy 1.17.1's brentq in one kernel, with the SPR target's log-likelihood.
15. A node's position as if it were the root (getLtqsAsIfRoot): one kernel.
16. nnnReorder shares the candidates it computes (bonsai_dataprocessing.py): round-robin over the processes instead of in contiguous blocks, with one allgather instead of three.
17. processOptReconfig tidies only what a move rebuilt (Tree.processOptReconfig): instead of two passes over the whole tree after every NNI move.
18. Throw-away trees (Tree.getOptEdgeList, nnnReorderRandom): freed by reference counting instead of the cyclic collector, and getOptEdgeList no longer repeats the whole-tree pass that optTimes ends with.
19. nnnReorderRandom keeps its node index between moves (bonsai_dataprocessing.py): a random candidate is found from per-node counts instead of rebuilding three lists before each move. The random draws are unchanged.
20. Loaded trees use plain-ndarray rows (load_data_for_tree): views of the same memory map, without np.memmap's Python hooks, below 25,000 cells.
21. Storing a tree (getLtqsComplete(mem_friendly=False)): one kernel per node for any number of children.
22. numba (new bonsai/bonsai_kernels.py, requirements.txt, requirements_bonsai_scout.txt): numba==0.68.0 is required. The kernels have explicit signatures, compile at the first import after an install, and cache=True keeps the result.

Every kernel repeats the numpy code it replaces operation by operation and in numpy's order, without fastmath, leaves logarithms to np.log, and hands back to the unchanged numpy code wherever numpy would warn; that is how the output stays byte-identical. bonsai/bonsai_lbfgsb.py and the brentq port in bonsai/bonsai_kernels.py are adapted from scipy 1.17.1 and carry scipy's BSD-3-Clause notice, and bonsai_lbfgsb.py also the MIT notice of scipy's _lbfgsb_py.py.

Tests and CI

The pytest suite from #99 passes on this branch: 843 passed, 4 skipped (bonsai_scout modules that need h5py). Bonsai has no CI workflows, so the suite was run in these environments:

Environment	Result
Python 3.14.3 with requirements.txt's pins, numba 0.68.0, mpi4py 4.1.2 on Open MPI 4.1.6	843 passed, 4 skipped
A fresh clone of main (b16ef6f) with this branch merged, a new environment from requirements-dev.txt	843 passed, 4 skipped

Python 3.12, an environment without mpi4py, the bonsai_scout environment, macOS and Windows were not run on this branch.

Behaviour changes

No output file changes. What a user can notice:

• numba is required: numba==0.68.0, with llvmlite 0.50.0. It installed without a build step here.
• The first run after an install compiles the kernels: a 16-cell run with an empty numba cache took 6.0 s against 3.1 and 3.5 s on main. Later runs load the compiled kernels.
• Cache files: numba writes 53 files to bonsai/__pycache__/, or to a per-user cache directory where that is not writable; NUMBA_CACHE_DIR overrides both.
• Warnings under MPI: merge pairs can now be evaluated on processes 1 and up, where Bonsai's warnings such as "Times cannot be optimized for this pair." are not printed and numpy's are printed by that process. Some nnnReorder progress lines come from another process.
• Pairs past a round's stop point: up to one fewer than the number of processes are evaluated and dropped. Such a pair could print a warning or reach the bare exit() in getOptTime2LeafTree; we did not see either on any input.
• No disp warning: Tree.optTimes no longer triggers scipy's DeprecationWarning for minimize's disp option, which Python hides by default.
• Other scipy versions: on a scipy other than 1.17.1, item 5 falls back to minimize and the speedup is smaller. The brentq port keeps 1.17.1's arithmetic with every scipy.

Usage benchmarks

Machine: AMD EPYC 9575F, 2 sockets x 64 cores (256 hardware threads, 2 NUMA nodes), 1 TB RAM, Ubuntu 24.04.5, Linux 7.0, Open MPI 4.1.6. Both versions are the checkouts and virtual environments of the fresh-clone check in How the output was checked: Python 3.14.3, numpy 2.4.6, scipy 1.17.1, pandas 3.0.3, scikit-learn 1.9.0 and mpi4py 4.1.2, plus numba 0.68.0 and llvmlite 0.50.0 for this branch.

Inputs: single-cell RNA-seq from 10x Genomics (CC BY 4.0), each put through Sanity 2.0.0 (-e 1 -v_m MAP) once, so both versions read the same Sanity output; and Bonsai's own example.

• Main case: 500 cells of pbmc_10k_v3, the input this change was developed on.
• Full experiments: the whole filtered matrices of pbmc_1k_v3, neuron_1k_v3, pbmc4k and pbmc_10k_v3, every cell and every gene with a non-zero count.
• Subset: 2,000 random cells of pbmc_10k_v3 (numpy.random.default_rng(20261002)). The 500- and 250-cell inputs are its first cells after default_rng(20261027).permutation(2000).
• Options and formats: the 250-cell input with zscore_cutoff -1, and the same input as features.txt and standard_deviations.txt (Sanity's delta.txt and d_delta.txt) with input_is_sanity_output False.
• Edge case: the 250-cell input plus exact copies of 50 of its cells (default_rng(7)), which end on zero-length branches.
• Tiny input: examples/example_data/simulated_binary_6_gens_samplingNoise (pre-2.0 Sanity layout), with nnn_n_randommoves 100 and nnn_n_randomtrees 2.

How it was timed:

• What runs: bonsai_main.py --step all from each checkout, under mpiexec --bind-to none -n <N> python -m mpi4py (plain python for one process), with a YAML from that checkout's create_config_file.py. Every option not named above is at the template's default.
• What the time includes: wall time from mpiexec starting to its exit: preprocessing, the reconstruction and the metadata step.
• Isolation: one configuration at a time, each version a fresh launch, with one BLAS and OpenMP thread per MPI process. The server is shared, and other jobs ran throughout: a launch waited, up to 10 minutes, while the load average was above 200, and the load average before the timed launches was between 33 and 187.
• Order: the version that ran first alternated from one input to the next.
• Compilation: each launch had its own new numba cache, filled by an untimed run of the 64-cell example at the same process count, so no timed run compiled.
• Peak memory: the maximum resident set of each MPI process (/usr/bin/time %M); the table gives the largest process.
• Output check: whether the seven files of the final folder were byte-identical (see How the output was checked).

Input class	Input	MPI processes	main (s)	This branch (s)	Speedup	main peak (MB)	This branch peak (MB)	Output
main case	PBMC, subset of pbmc_10k_v3, 500 cells, 16,758 genes	1	167.54	32.56	5.14x	275	275	identical
main case	PBMC, subset of pbmc_10k_v3, 500 cells, 16,758 genes	2	109.73	19.59	5.60x	268	273	identical
main case	PBMC, subset of pbmc_10k_v3, 500 cells, 16,758 genes	4	79.22	13.03	6.08x	271	275	identical
main case	PBMC, subset of pbmc_10k_v3, 500 cells, 16,758 genes	8	76.16	12.28	6.20x	271	277	identical
main case	PBMC, subset of pbmc_10k_v3, 500 cells, 16,758 genes	16	61.04	11.23	5.44x	261	278	identical
full experiments	pbmc_1k_v3, human PBMC, 1,222 cells, 18,391 genes	1	428.93	81.01	5.29x	341	340	identical
full experiments	pbmc_1k_v3, human PBMC, 1,222 cells, 18,391 genes	4	265.14	45.63	5.81x	346	344	identical
full experiments	pbmc_1k_v3, human PBMC, 1,222 cells, 18,391 genes	16	185.34	33.87	5.47x	328	353	identical
full experiments	neuron_1k_v3, E18 mouse brain, 1,301 cells, 18,085 genes	1	613.03	127.06	4.82x	502	456	identical
full experiments	neuron_1k_v3, E18 mouse brain, 1,301 cells, 18,085 genes	4	329.37	49.54	6.65x	522	460	identical
full experiments	neuron_1k_v3, E18 mouse brain, 1,301 cells, 18,085 genes	16	222.56	47.03	4.73x	465	460	identical
full experiments	pbmc4k, human PBMC, 4,340 cells, 19,773 genes	4	1,010.64	128.44	7.87x	392	453	identical
full experiments	pbmc4k, human PBMC, 4,340 cells, 19,773 genes	16	767.36	105.02	7.31x	375	449	identical
full experiments	pbmc_10k_v3, human PBMC, 11,769 cells, 23,036 genes	16	3,526.12	525.62	6.71x	1,083	1,307	identical
subset	PBMC, subset of pbmc_10k_v3, 2,000 cells, 19,671 genes	4	377.81	56.95	6.63x	397	414	identical
options and formats	PBMC subset, 250 cells, all genes (zscore_cutoff -1), 15,143 genes	4	65.89	13.03	5.06x	357	281	identical
options and formats	PBMC subset, 250 cells, as features.txt and standard_deviations.txt, 15,143 genes	4	45.77	10.78	4.25x	230	250	identical
edge case	PBMC subset, 300 cells, 50 of them copies, 15,143 genes	4	53.53	9.48	5.65x	256	260	identical
tiny input	Bonsai's simulated example, 64 cells, 12,304 genes	1	9.98	3.12	3.20x	192	246	identical
tiny input	Bonsai's simulated example, 64 cells, 12,304 genes	4	10.03	2.77	3.62x	194	248	identical

The output files were byte-identical in every configuration.

Over the 9 configurations of the four full experiments, the geometric mean speedup is 5.99x (4.73x to 7.87x), and the summed wall time is 7,348 s for main and 1,143 s for this branch.

How the output was checked

The usage benchmarks above check speed and memory. Separately, we compared main and this branch as follows:

• Scale: 77 comparisons in five checks, from 6 to 11,769 cells, as one plain process and at 1 to 16 MPI processes.
• What was compared: the seven files of the final folder (tree.nwk, edgeInfo.txt, vertInfo.txt, metadata.json, metadata_bonsai.txt and the two posterior_*_vertByGene.npy), byte for byte. The absolute results and data paths in metadata.json were replaced by tokens first, and the reconsTime column, the run's wall time, was removed from metadata_bonsai.txt.

What ran

Check	Comparisons	What
Small cases	47	18 whole reconstructions of 6 to 128 cells at 0 to 4 MPI processes (Sanity 2.0 MAP and MLE and pre-2.0 output, features with and without standard deviations, duplicated cells, all genes, no upper bounds, use_knn 3, the long_branch_sorted and super_sure SPR strategies, one signal gene), and 29 direct calls of the changed functions on small and unusual inputs
Main case	3	the 500-cell input at 4 processes, in three orders of its cells
Validation	1	the 2,000-cell subset at 4 processes
Fresh clone	6	upstream main (b16ef6f) cloned anew with this branch merged, each version in a new environment from its own requirements-dev.txt; the main case in three cell orders, each twice, alternating which version went first
Usage benchmarks	20	every configuration of the table above

Results

Every comparison was identical. The usage benchmarks' 2,000-cell outputs also differed, in both versions, from the validation row's main output in one field: metadata.json's dataset, which records the dataset name the YAML gives, and the two runs gave the dataset different names. A small case that aborted was retried up to three times, because main itself aborts now and then with spr_strategy: super_sure under MPI. At one process, where items 3, 8 and 16 do nothing, the speedup is still 5.14x on the main case, 5.29x on pbmc_1k_v3 and 4.82x on neuron_1k_v3, and both versions use 0.97 to 1.00 CPU-seconds per wall second, so the gain does not come from extra threads.

Memory

Peak resident memory of the largest MPI process ranges from 0.79 to 1.28 times main's, with a median of 1.02. It is lower on the all-genes input (281 MB against 357 MB) and on neuron_1k_v3 at 1 and 4 processes (456 to 460 MB against 502 to 522 MB), and higher on the 64-cell example (246 to 248 MB against 192 to 194 MB), pbmc4k (449 to 453 MB against 375 to 392 MB) and pbmc_10k_v3 at 16 processes (1,307 MB against 1,083 MB). numba and the compiled kernels take about 46 MB per process, most of the difference on the 64-cell example, and Tree.optTimes holds a table of four to five float64 values per node and gene while it runs. We did not measure how much of the increase on pbmc4k and pbmc_10k_v3 comes from each.

What this testing does not cover

• Platforms: only x86-64 Linux on one AMD Zen 5 CPU, with Python 3.14.3, numpy 2.4.6, scipy 1.17.1 and Open MPI 4.1.6. We ran nothing on macOS, Windows, ARM, another Python or other library versions.
• Size: no input above 11,769 cells, so not the mem_friendly mode Bonsai switches on above 25,000 cells, where this branch keeps main's code in nnnReorder, load_data_for_tree and Tree.optTimes.
• Steps and options: the separate --step runs and pickup_intermediate were not compared, though the test suite runs them and passes; other options were compared only in the small cases.
• Other entry points: backbone_based_bonsai/, the scripts that add cells, and bonsai_scout were checked only by the test suite.
• Logs: not compared line by line. Both versions printed the same number of Bonsai's "Watch out, the number of nodes stored in tree.nNodes is not equal to the number I just counted." warning on every usage-benchmark input, and neither printed an error.

Bugs in main found along the way

While writing the tests and checking this change, we reproduced bugs in main and reported them in #98. This pull request fixes none of them, since each fix would change some output. Two touch this change: the occasional abort of spr_strategy: super_sure under MPI, and the bare exit() in getOptTime2LeafTree (see Behaviour changes).

🤖 Generated with Claude Code

[jmschrei]

Hello! The above message was written by Claude, but this comment is written by me (though based on my other auto-optimization PRs).

I have been developing a tool, called [auto-optimize], which basically applies the auto-research loop to code to improve speed and memory usage of code, but with extensive robustness checks to try to ensure nothing subtly breaks. I have been applying [auto-optimize] across the genomics ecosystem, starting with my own tools but now extending to other important ecosystem tools, with the hope of making things broadly better. Most recently, I've applied it to MACS3, Bowtie2, chromap, Picard MarkDuplicates, and bedtools.

In an effort to be more robust, the process first augments the unit tests to make sure that nothing breaks during the search. See #99 (feel free to merge that independently of any choice you make here). Any iteration in our loop that broke any unit test was automatically discarded, and we used a set of external validations that were not directly optimized against each iteration to make sure we got identical results. At the end, we ran Bonsai several times to make sure it still reproduced on real inputs.

If you look at the results below, you'll see that there were several big improvements, but also a collection of small improvements that all collectively led to the improved speed. (ignore the orange line)

In an effort to be consistent with https://rewrites.bio/, I'd like this to be merged into this repository so the original developers and maintainers can continue to get credit for their work.

However, as a maintainer myself, I understand being skeptical of merging a PR of this size, particularly when the submitter can't explain everything in it (a consequence of so many changes being made at a low level). There are several paths forward: (1) you can merge the PR as-is, (2) you can push your own commits on top of it to revert any changes you dislike or built upon them, (3) you can cherry-pick out any individual or set of improvements you like and merge them as separate PRs (as Ben is doing with bowtie2 right now), (4) you can ignore this entirely.

We're happy to share any logs or information that would help you!

[dhdegroot]

For now, we will not be adding many lines of AI-generated code to this repo, at least not if I don't have time to check it carefully.