tem1-bo: derivative-free optimizers on a protein fitness landscape

The problem

We want a TEM-1 beta-lactamase variant with high fitness (antibiotic resistance, Gray et al. 2018 / PEER benchmark; the train/valid/test splits are merged into 5198 single mutants). Each variant is represented by the mean-pooled final-layer embedding of ESM2-8M (320 dims). Those embeddings are pre-computed and shipped in src/tem1_bo/data/tem1.npz (6 MB), so nothing is downloaded and no GPU / torch is needed.

To get a continuous black box, the first d principal components of the embeddings (default d = 8) span a box, each axis rescaled to [0, 1]:

• x in [0, 1]^d: a point in embedding space
• f(x): minus the measured fitness of the variant whose embedding is nearest to x (everything is minimized)

f is piecewise constant (one plateau per variant), so it has no useful gradient. It is deterministic (not noisy) and one evaluation takes about 0.1 ms here. The real-world analogue is a lab measurement of a variant, which is why sample efficiency is the interesting question.

Optimizers

All get the same number of evaluations T per seed.

Run

uv sync
uv run q2 --help
uv run q2 compare --budget 100 --seeds 10 --out results/
uv run pytest

Without uv (plain venv + pip)

Python >= 3.11. From this folder:

python3 -m venv .venv
source .venv/bin/activate          # Windows: .venv\Scripts\activate
pip install .                      # builds and installs the package and q2
q2 --help
q2 compare --budget 100 --seeds 10 --out results/

For development use pip install -e . pytest and run pytest. To build a wheel instead: pip install build && python -m build (writes dist/), then pip install dist/*.whl in any other venv. pip install resolves the dependencies from pyproject.toml (uv.lock is only used by uv), so versions may differ slightly from the locked ones.

The defaults (--budget 100 --seeds 10) take about 2 min on a laptop CPU. Other options: --dim (box dimension), --optimizers (subset). Seeds 0 .. S-1 are used for everything, running twice gives identical results.json.

Output

• results/results.json: problem description, budget and one entry per (optimizer, seed) with the value f of every evaluation in order (x has the evaluated points).
• results/ranking.json: for every (optimizer, seed) the Spearman rho between the true f of all 5198 variants and the prediction of a GP fitted on the first t evaluations (t = 10, 20, ..., budget). The same surrogate is used for all optimizers, so it measures how informative each optimizer's evaluations are about the whole landscape. If all evaluated f are equal the prediction is constant and rho is set to 0.
• results/comparison.png, three panels:
  1. best f found so far against the number of evaluations (mean over seeds, band = +-1 std; lower is better, f = -fitness)
  2. histogram of f over all 5198 variants in the dataset (same y axis as panel 1, log counts)
  3. the Spearman rho curves from ranking.json (mean +- std over seeds)