Metadata-Version: 2.4
Name: nucleobench
Version: 2.0.10
Summary: Nucleobench optimizers and tasks.
Author-email: Move37 Labs <joel.shor@move37labs.com>
Project-URL: Homepage, https://github.com/move37-labs/nucleobench
Project-URL: Bug Tracker, https://github.com/move37-labs/nucleobench/issues
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1
Requires-Dist: lightning
Requires-Dist: pytorch-lightning<3.0,>=2.0
Requires-Dist: gReLU==1.1.0
Requires-Dist: bpnet-lite==0.8.1
Requires-Dist: numpy>=2.0.0
Requires-Dist: scipy
Requires-Dist: pandas
Requires-Dist: h5py
Requires-Dist: pyarrow
Requires-Dist: fastparquet
Requires-Dist: google-cloud-storage
Requires-Dist: gcsfs
Requires-Dist: pyBigWig==0.3.25
Requires-Dist: xxhash
Requires-Dist: ml-collections
Requires-Dist: tqdm
Requires-Dist: einops
Requires-Dist: gradabeam
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: pytest-xdist; extra == "dev"
Requires-Dist: pytest-repeat; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: coverage; extra == "dev"
Requires-Dist: pyclean; extra == "dev"
Dynamic: license-file

## NucleoBench

[![codecov](https://codecov.io/gh/move37-labs/nucleobench/graph/badge.svg?token=8ZVWDF7BT3)](https://codecov.io/gh/move37-labs/nucleobench)
[![PyPI version](https://badge.fury.io/py/nucleobench.svg)](https://badge.fury.io/py/nucleobench)
[![Python Versions](https://img.shields.io/pypi/pyversions/nucleobench.svg)](https://pypi.org/project/nucleobench/)
[![License](https://img.shields.io/github/license/move37-labs/nucleobench)](https://github.com/move37-labs/nucleobench/blob/main/LICENSE)
[![DOI](https://img.shields.io/badge/bioRxiv-10.1101%2F2025.06.20.660785-b31b1b.svg)](https://www.biorxiv.org/content/10.1101/2025.06.20.660785)
[![Conda Environment](https://img.shields.io/badge/conda-env-green)](https://github.com/move37-labs/nucleobench/blob/main/environment.yml)
[![Docker Hub](https://img.shields.io/badge/Docker%20Hub-nucleobench-blue?logo=docker)](https://hub.docker.com/r/joelshor/nucleobench)
[![GHCR](https://img.shields.io/badge/ghcr.io-nucleobench-blue?logo=github)](https://github.com/move37-labs/nucleobench/pkgs/container/nucleobench)

| Environment  | Unit Tests | Docker Integration Tests |
| :--- | :--- | :--- |
| **Micromamba** | [![Unit Tests (Micromamba)](https://github.com/move37-labs/nucleobench/actions/workflows/unit-tests-in-micromamba.yml/badge.svg)](https://github.com/move37-labs/nucleobench/actions/workflows/unit-tests-in-micromamba.yml) | [![Docker Integration Tests (Micromamba)](https://github.com/move37-labs/nucleobench/actions/workflows/docker-integration-tests-in-micromamba.yml/badge.svg)](https://github.com/move37-labs/nucleobench/actions/workflows/docker-integration-tests-in-micromamba.yml) |
| **GHCR** | [![Unit Tests (GHCR)](https://github.com/move37-labs/nucleobench/actions/workflows/unit-tests-in-docker-github.yml/badge.svg)](https://github.com/move37-labs/nucleobench/actions/workflows/unit-tests-in-docker-github.yml) | [![Docker Integration Tests (GHCR)](https://github.com/move37-labs/nucleobench/actions/workflows/docker-integration-tests-in-docker-github.yml/badge.svg)](https://github.com/move37-labs/nucleobench/actions/workflows/docker-integration-tests-in-docker-github.yml) |

### Model Integration Tests (Micromamba)
[![Enformer Integration Tests (Micromamba)](https://github.com/move37-labs/nucleobench/actions/workflows/enformer-integration-tests-in-micromamba.yml/badge.svg)](https://github.com/move37-labs/nucleobench/actions/workflows/enformer-integration-tests-in-micromamba.yml)

**A large-scale benchmark for modern nucleic acid sequence design algorithms (NucleoBench), and a new design algorithm that outperforms existing designers (GrAdaBeam).  Link to ICML GenBio 2025 workshop paper [here](https://www.biorxiv.org/content/10.1101/2025.06.20.660785).**

[comment]: <> (Consider an image here.)

This repo is intended to be used in a few ways:
1. Design a DNA sequence with selective expression in a cell-type (or any other target property in the benchmark, see list [here](#summary-of-tasks-in-nucleobench)), using the GrAdaBeam algorithm (or any of the ones listed [here](#summary-of-designers-in-nucleobench))
2. Design a DNA sequence with high binding affinity for a specific transcription factor (such as the ones listed [here](#summary-of-tasks-in-nucleobench)), using the GrAdaBeam algorithm (or any of the ones listed [here](#summary-of-designers-in-nucleobench))
1. Design a DNA or RNA sequence for a new task, using any designer (see tutorial [here](recipes/colab/custom_task.ipynb))
1. Run a new design algorithm on NucleoBench tasks.
1. Reproduce the NucleoBench results, using the standard tasks / designers or using your custom ones, on the Cloud using Google Batch or AWS (instructions [here](#get-started-in-8-minutes-google-batch-inference))


### Citation

Please cite the following publication when referencing NucleoBench or GrAdaBeam:

```
@article{shor2025nucleobench,
  author    = {Shor, Joel and Strand, Erik and McLean, Cory Y.},
  title     = {{GrAdaBeam: Combining model gradients with evolutionary search for generalizable nucleic acid design}},
  journal   = {bioRxiv},
  year      = {2026},
  doi       = {10.1101/2025.06.20.660785},
  url       = {https://www.biorxiv.org/content/10.1101/2025.06.20.660785}
}
```

## Contents

- [Quick Start](#quick_start)
  - [1 minute install w/ pip](#get-started-in-1-minute-pip-install)
  - [5 minute install w/ source](#get-started-in-5-minutes-git-clone)
  - [8 minute install w/ Google Batch inference](#get-started-in-8-minutes-google-batch-inference)
- [Details](#details)
  - [GrAdaBeam](#gradabeam)
- [FAQ](#faq)

## Quick Start

NucleoBench is provided via **PyPi** or **source**.

### Get started in 1 minute (pip install)

Requires Python 3.10+. On Debian/Ubuntu, install build deps for `pyBigWig` first:
```bash
sudo apt-get install -y libcurl4-openssl-dev
```

Then install `nucleobench` on your terminal:
```bash
# Standard / GPU install:
pip install nucleobench
# CPU-only install (forces PyTorch to use CPU version):
pip install nucleobench --extra-index-url https://download.pytorch.org/whl/cpu
```

This installs PyTorch and related scientific deps and can take several minutes.

Then run in Python:
```python
# 1. Choose a model (task).
from nucleobench import models
model = models.get_model('substring_count')
model_init_args = model.debug_init_args()
model_init_args['substring'] = 'ATGTC'
model_fn = model(**model_init_args)

# 2. Choose an optimizer.
from nucleobench import optimizations
opt_obj = optimizations.get_optimization('gradabeam')
opt_init_args = opt_obj.debug_init_args()
opt_init_args['model_fn'] = model_fn
opt_init_args['start_sequence'] = 'A' * 100
designer = opt_obj(**opt_init_args)

# 3. Run the designer and show the results.
designer.run(n_steps=100)
ret = designer.get_samples(1)
ret_score = model_fn(ret)
print(f'Final score: {ret_score[0]}')
print(f'Final sequence: {ret[0]}')
```

Output:
```bash
Final score: -535.0
Final sequence: ATGTCTGTCTATGTCATGTCATGTCATGTCATGTCATGTCTCTATGTCATGTCTATGTCTATGTCTATGTCATGTCATGTCTGTCTATGTCATGTATGTC
```

This "recipe" can be found under [`recipes/python/gradabeam_substringcount.py`](recipes/python/gradabeam_substringcount.py).

### Get started in 5 minutes (git clone)

```bash
git clone https://github.com/move37-labs/nucleobench.git
cd nucleobench
conda env create -f environment.yml
conda activate nucleobench
```

Now run the main entrypoint:

```bash
python -m docker_entrypoint \
    --model substring_count \
        --substring 'ATGTC' \
    --optimization gradabeam \
        --beam_size 2 \
        --n_rollouts_per_root 4 \
        --mutations_per_sequence 2 \
        --exploration_alpha 0.05 \
        --rng_seed 0 \
    --max_seconds 15 \
    --optimization_steps_per_output 5 \
    --proposals_per_round 2 \
    --output_path ./output/python_recipe/gradabeam_substringcount \
    --start_sequence AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
```

Output:
```bash
...
Completed round 1839 (5 steps) took 0.01s. Avg 0.00s per step.
  0%|                                     | 1840/99999999 [00:14<226:20:41, 122.72it/s]
Proposals deposited at:
    ./output/python_recipe/gradabeam_substringcount/gradabeam_substring_count/20260725_180740/20260725_180755.parquet
```

This "recipe" can be found under [`recipes/python/gradabeam_substringcount.sh`](recipes/python/gradabeam_substringcount.sh).

### Get started in 8 minutes (Google Batch inference)

Google Batch is Google's cheapest batch compute offering. It enables relatively cheap parallel compute on the cloud.

First, setup a Google Cloud project by following instructions [here](https://cloud.google.com/batch). You will need to activate the Google Batch API. If you want the job to write to a private bucket, you will also need to setup what's called a "Service Account" with the proper permissions. Once you have done this, you will need to collect the following information from your new project, and fill this information in [runners/gcp/config.py](runners/gcp/config.py):

1. PROJECT_ID (the project ID of the project you created above)
1. REGION (the region of the project)
1. BUCKET_NAME (the bucket of the output)
1. SERVICE_ACCOUNT_EMAIL (option: a service account for security)

Next, create the conda environment to launch jobs to the Cloud runner:
```bash
conda env create -f runners/environment.yml
```
Output:
```bash
Transaction finished

To activate this environment, use:

    conda activate runners

Or to execute a single command in this environment, use:

    conda run -n runners mycommand
```

Now run the dry run Google Batch job with the test script:

```bash
conda activate runners
python -m runners.gcp.job_launcher \
    --dry-run \
    --verbose \
    runners/testdata/adabeam_test.tsv
```

Output:
```bash
2025-08-28 18:00:21,651 - INFO - Row 1: Generated job name: bpnet-rad21-adabeam-001
2025-08-28 18:00:21,651 - INFO - Loaded 1 jobs from runners/testdata/adabeam_test.tsv
2025-08-28 18:00:21,651 - INFO - DRY RUN MODE - No jobs will be launched
2025-08-28 18:00:21,651 - INFO - Would launch job: bpnet-rad21-adabeam-00120250828-18-00-21
```

Now run the real job:
```bash
python -m runners.gcp.job_launcher \
    --verbose \
    runners/testdata/adabeam_test.tsv
```

Output:
```bash
2025-08-28 18:01:32,909 - INFO - Row 1: Generated job name: bpnet-rad21-adabeam-001
2025-08-28 18:01:32,909 - INFO - Loaded 1 jobs from runners/testdata/adabeam_test.tsv
2025-08-28 18:01:32,910 - DEBUG - Checking None for explicit credentials as part of auth process...
2025-08-28 18:01:32,910 - DEBUG - Checking Cloud SDK credentials as part of auth process...
2025-08-28 18:01:33,528 - DEBUG - Starting new HTTPS connection (1): oauth2.googleapis.com:443
2025-08-28 18:01:33,708 - DEBUG - https://oauth2.googleapis.com:443 "POST /token HTTP/1.1" 200 None
2025-08-28 18:01:34,087 - INFO - Successfully launched job: projects/nucleorave/locations/us-central1/jobs/bpnet-rad21-adabeam-00120250828-18-01-32
2025-08-28 18:01:34,087 - INFO - 
Job Launch Summary:
2025-08-28 18:01:34,087 - INFO - Successful: 1
2025-08-28 18:01:34,087 - INFO - Failed: 0
2025-08-28 18:01:34,087 - INFO - Successful jobs: bpnet-rad21-adabeam-00120250828-18-01-32
```

Finally, make sure that the job completes fully, and the output is in the right bucket.

## Details

We introduce **GrAdaBeam**, a hybrid model-based optimization algorithm that
combines gradient-derived attention maps with an adaptive beam search to navigate complex nucleic acid fitness landscapes. By unifying the broad exploration of evolutionary methods with the precise guidance of gradient descent, GrAdaBeam overcomes a central limitation of existing approaches: no single optimization strategy performs robustly across the full spectrum of genomic design tasks. We rigorously evaluate GrAdaBeam and nine other design algorithms using NucleoBench, a novel benchmark covering 17 diverse genomic tasks that introduces a paired-start-sequence design for superior statistical comparisons. GrAdaBeam statistically outperforms all other algorithms across over 600,000 experiments, never ranking lower than second across all 17 benchmark tasks, while baseline methods often struggle on large models or long sequences. Critically, GrAdaBeam sequences generalize most reliably to independent predictive models and recover canonical transcription factor binding motifs de novo, providing evidence of biological signal capture beyond the optimization target. GrAdaBeam and NucleoBench are freely available as an open-source package.

### GrAdaBeam

GrAdaBeam is a gradient-guided adaptive beam search. At each step it
samples how many sites to edit, chooses those sites with a mixture of
uniform exploration and oracle-gradient (TISM) signal, proposes new
bases, and keeps improving children. Mutation rate and the exploration
weight are updated from successful edits so the search adapts during
optimization.

<div align="center">
<img src="assets/images/gradabeam_flowchart.png" alt="GrAdaBeam algorithm schematic" style="width: 70%; max-width: 800px; height: auto;" />
<br/>
<small>One GrAdaBeam step: the oracle supplies gradients for position selection (TISM) and forward scores for acceptance; accepted children update the parent sequence and the adaptive parameters μ and α.</small>
</div>

<div align="center">
<img src="assets/images/order_score_violin.png" alt="Order-score violin plot: GrAdaBeam ranks highest across NucleoBench tasks" style="width: 70%; max-width: 800px; height: auto;" />
<br/>
<small>Order scores across NucleoBench tasks (higher is better). For a given task and start sequence, the order score is the number of designers with worse final fitness.</small>
</div>

### Summary of tasks in NucleoBench

| TASK CATEGORY | MODEL | DESCRIPTION | NUM TASKS | SEQ LEN (BP) | INFERENCE (MS / EXAMPLE) |
| :--- | :--- | :--- | :--- | :--- | :--- |
| Cell-type specific cis-regulatory activity | Malinois | How DNA sequences control gene expression from the same DNA molecule. Cell types are: *precursor blood cells*, *liver cells*, *neuronal cells*. | 3 | 200 | 2 |
| Transcription factor binding | BPNet-lite | How likely a specific transcription factor (TF) will bind to a particular stretch of DNA. Specific TFs: *CTCF*, *E2F3*, *ELF4*, *GATA2*, *JUNB*, *MAX*, *MECOM*, *MYC*, *OTX1*, *RAD21*, *SOX6* | 11 | 3000 | 14 |
| Chromatin accessibility | BPNet-lite | How physically accessible DNA is for interactions with other molecules. | 1 | 3000 | 120 |
| Ribosomal loading | RiNALMo | Mean ribosome load on 5' UTR sequences (translational efficiency). | 1 | 100 | 450 |
| Selective gene expression | Enformer | Prediction of gene expression. We optimize for *maximal expression in muscle cells, minimal expression in liver cells*. | 1 | 196,608 / 256 * | 7900 |

<small>*Input length is 196,608 bp, but only 256 bp are edited.</small>

### Summary of designers in NucleoBench

| Algo | Description | Gradient-based |
| :--- | :--- | :---: |
| Directed Evolution | Random mutations, track the best. | ❌ |
| Simulated Annealing | Greedy optimization with random jumps. | ❌ |
| [AdaLead](https://arxiv.org/abs/2010.02141) | Heuristic-guided random-mutation search. | ❌ |
| [FastSeqProp](https://bmcbioinformatics.biomedcentral.com/articles/10.1186/s12859-021-04437-5) | Straight-through estimator for maximal input. | ✅ |
| [Ledidi](https://www.biorxiv.org/content/10.1101/2020.05.21.109686v1) | Gumbel softmax sampling for maximal input. | ✅ |
| GFlowNet* | Flow-based generative model. | ❌ |
| --- |
| Ordered Beam | Greedy beam search, in fixed edit order. | ❌ |
| Unordered Beam | Greedy beam search, any edit order. | ❌ |
| Gradient Evo | Directed Evolution, guided by gradients. | ✅ |
| [GrAdaBeam](https://www.biorxiv.org/content/10.1101/2025.06.20.660785) | Efficient beam search guided by gradients. | ✅ |

<small>Table: Summary of designers in NucleoBench. Above the solid line are designers already found in the nucleic acid design literature.
Below the line are designers from the search literature not previously used to benchmark nucleic acid sequence design and hybrid
algorithms devised in this work. *GFlowNet is a generative model, not a model-based optimizer, included for benchmark completeness.
GFlowNets do not use gradients of the oracle model.</small>

## FAQ

1. How can I add a new task to NucleoBench?
    A: Follow [this](recipes/colab/custom_task.ipynb) colab notebook.
