Metadata-Version: 2.5
Name: tem1-bo
Version: 0.1.0
Summary: Derivative-free optimizers on the TEM-1 beta-lactamase fitness landscape (ESM2 embedding space)
Project-URL: Repository, https://github.com/j-albr16/ml
Author: Jan Albrecht
License-Expression: MIT
License-File: LICENSE
Keywords: bayesian-optimization,cma-es,derivative-free-optimization,esm2,protein-fitness
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Requires-Python: >=3.11
Requires-Dist: cma>=3.3
Requires-Dist: matplotlib>=3.8
Requires-Dist: numpy>=1.26
Requires-Dist: scipy>=1.11
Description-Content-Type: text/markdown

# 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.

| name | family | implementation |
|---|---|---|
| `random` | baseline | uniform samples |
| `sobol` | space-filling | scrambled Sobol (`scipy.stats.qmc`) |
| `nelder-mead` | direct search | `scipy.optimize`, restarted from a random point when it converges |
| `cma-es` | direct search / ES | `cma`, sigma0 = 0.3, restarted when it stops |
| `bo` | model-based | own numpy GP (RBF, hypers by marginal likelihood) + lower confidence bound (beta = 2), 10 Sobol initial points |

## Run

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

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)
