Sampleworks

This repository is under active development. Please always use the latest version. If you encounter any problems, please create an issue on GitHub and include: the PDB ID, the CIF file you used, your density map(s), and log information.

We would welcome contributions from the community. We are most interested in:

• new ModelWrappers for additional structure prediction models (especially smaller models which may be more steerable)
• fast, differentiable modules to allow guidance from other experimental data modalities besides X-ray electron density.

Sampleworks is a Python framework for integrating generative biomolecular structure models with experimental data. Read our blog post for an introduction.

Why sampleworks?

Biomolecular structure prediction and design models are currently trained on single state structures and fail to accurately predict the ensemble of conformations each macromolecule occupies. But there is still hope! Current models show promise in capturing the underlying distribution of realistic macromolecular structures. We want to utilize the prior represented in these models and experimental observations to improve the sampling of the underlying ensemble present in the experiment and use this information to both understand biomolecular function and improve ensemble prediction.

Currently, each structure prediction model has a different implementation, requiring bespoke boilerplate code to plug each model into experimental guidance. Our goal is to resolve this and expand the experimental methods we can provide guidance with. This will open new opportunities for model evaluation directly against experimental data, and help unlock new sources of data for training the next generation of biomolecular structure predictors.

Installation

Requirements: Linux x86-64, CUDA 12, Python ≥ 3.11, < 3.14

1. Install Pixi

curl -fsSL https://pixi.sh/install.sh | sh

2. Clone and install

git clone git@github.com:diff-use/sampleworks.git
cd sampleworks
pixi install -a   # install all environments

Note: pixi install -a resolves all environments. This (currently) requires CUDA 12 and will fail on machines without it.

Each generative model has its own Pixi environment. Install only what you need:

pixi install -e boltz      # Boltz-1 / Boltz-2
pixi install -e protenix   # Protenix
pixi install -e rf3        # RosettaFold3

3. Download model checkpoints

Boltz-1 and Boltz-2 (stored in ~/.boltz/):

pixi run -e boltz python -c "
from boltz.main import download_boltz1, download_boltz2
import pathlib
cache = pathlib.Path('~/.boltz/').expanduser()
download_boltz1(cache)
download_boltz2(cache)
"

Protenix: checkpoint is downloaded automatically on first use.

RosettaFold3 (RF3): see the RC-Foundry repository for instructions. Default path: ~/.foundry/checkpoints/rf3_foundry_01_24_latest.ckpt

Quick Start

Run Boltz-2 pure guidance on the included 1VME example:

pixi run -e boltz sampleworks-guidance \
    --model boltz2 \
    --guidance-type pure_guidance \
    --protein 1VME \
    --model-checkpoint ~/.boltz/boltz2_conf.ckpt \
    --structure tests/resources/1vme/1vme_final_carved_edited_0.5occA_0.5occB.cif \
    --density tests/resources/1vme/1vme_final_carved_edited_0.5occA_0.5occB_1.80A.ccp4 \
    --resolution 1.8 \
    --output-dir output/boltz2_pure_guidance \
    --guidance-start 130 \
    --ensemble-size 4 \
    --augmentation \
    --align-to-input

Output files appear in output/boltz2_pure_guidance/: refined.cif (final ensemble), losses.txt, trajectory/, run.log.

CLI reference

sampleworks-guidance is the unified command-line interface for running guidance on a single structure.

Required arguments:

Argument	Description
--model	boltz1, boltz2, protenix, or rf3
--guidance-type	pure_guidance or fk_steering
--protein	Protein identifier (should match naming used in grid search / evaluation)
--structure	Path to input structure file (CIF)
--density	Path to density map (CCP4/MRC/MAP)
--resolution	Map resolution in Angstroms

Model-specific arguments (e.g. --method for boltz2, --msa-path for rf3) and guidance-type-specific arguments (e.g. --num-particles for fk_steering) are included automatically. Run sampleworks-guidance --model <model> --guidance-type <type> --help to see all available options.

Grid Search

run_grid_search.py sweeps a model across scalers, ensemble sizes, and gradient weights:

pixi run -e boltz python run_grid_search.py \
    --proteins proteins.csv \
    --models boltz2 \                # options: boltz1, boltz2, protenix, rf3 (make sure env aligns!)
    --methods "X-RAY DIFFRACTION" \  # only useful for Boltz-2, ignored otherwise
    --scalers pure_guidance \        # options: pure_guidance, fk_steering, or both as space-separated list
    --ensemble-sizes "1 4" \
    --gradient-weights "0.1 0.2" \
    --output-dir grid_search_results \
    --gradient-normalization \       # normalize guidance update magnitude to diffusion update magnitude
    --augmentation \                 # apply random rotations and translations at each step (defaults for inference with AF3-like models)
    --align-to-input                 # align to input structure at each step (required for density guidance to work since it is not rotation/translation invariant)

proteins.csv format

Required columns and format. Supported density map formats: .ccp4, .mrc, .map (not MTZ or SF-CIF yet).

name,structure,density,resolution
1abc,/data/structures/1abc.cif,/data/maps/1abc.ccp4,2.0
2xyz,/data/structures/2xyz.cif,/data/maps/2xyz.mrc,1.8

Key arguments:

Argument	Description	Default
--proteins	CSV with structure/density/resolution columns	required
--models	Model to run. One of boltz1, boltz2, protenix, rf3	required
--scalers	Guidance method(s) to sweep	pure_guidance fk_steering
--ensemble-sizes	Space-separated values, e.g. "1 4"	"1 2 4 8"
--gradient-weights	Space-separated values, e.g. "0.1 0.2"	"0.01 0.1 0.2"
--methods	Boltz-2 sampling method (required for boltz2)	X-RAY DIFFRACTION
--max-parallel	Parallel workers (default: number of GPUs)	auto
--dry-run	Print jobs without running them	off
--force-all	Re-run including already-successful jobs	off
--only-failed	Re-run only failed jobs	off
--only-missing	Run only jobs not yet started	off

Output layout: grid_search_results/<protein>/<model>[_<method>]/<scaler>/ens<N>_gw<W>/

Note: Jobs are skipped if a refined.cif file already exists in the output directory. Some flags (e.g., --use-tweedie, --gradient-normalization) are not reflected in the directory structure, so changing them alone won't trigger a re-run. Use --force-all to re-run all jobs regardless. This is under active development and will likely change soon.

Instructions for running evaluation and metrics scripts are coming soon.

Running preset experiments on ACTL (run_experiments)

This section is Astera-specific: it assumes access to ACTL, the internal Harbor image registry, and the diffuse-shared PVC. External users can run the same TOML presets with sampleworks-runs or python -m sampleworks.runs.cli after setting equivalent local paths for DATA_DIR, PROTEINS_CSV, RESULTS_DIR, MSA_CACHE_DIR, and model checkpoints.

Start an 8-GPU ACTL machine named sampleworks with the Sampleworks image and the shared data volume mounted:

actl pod up sampleworks --profile 8x --image harbor.astera.sh/library/pixi-with-checkpoints:sampleworks --storage shared --pvc-size 200Gi --mount diffuse-shared --yes

Keep that terminal open; it maintains sync and SSH. From another terminal:

actl pod status sampleworks
# copy the `ssh:` line, then run it, for example:
ssh workspace.actl-ws-<user>-sampleworks.devspace
cd /home/dev/workspace

The main command is run_experiments. It reads TOML presets and launches the right run_grid_search.py jobs, pixi environments, GPU assignments, logs, results directory, and MSA cache.

export DATA_DIR=/mnt/diffuse-shared/raw/sampleworks/initial_dataset_40_occ_sweeps
export PROTEINS_CSV="$DATA_DIR/proteins.csv"
export SAMPLEWORKS_ACTL_RUN_NAME="$(hostname -s)"

run_experiments --list        # show available presets (does not require DATA_DIR)
run_experiments --show rf3    # inspect what will run
run_experiments --dry-run rf3 # print commands without running
run_experiments rf3           # run the standalone RF3 preset
run_experiments boltz         # run Boltz2 X-ray + Boltz2 MD
run_experiments boltz1        # run standalone Boltz1
run_experiments protenix      # run the standalone Protenix preset
run_experiments full_8gpu     # run the full 8-GPU comparison preset

The default full_8gpu preset runs Boltz2 XRD, Boltz2 MD, RF3, and Protenix in parallel. Run a subset with:

run_experiments full_8gpu --jobs rf3,protenix

Standalone presets are available for each model/model family: boltz, boltz1, boltz2, boltz2_xrd, boltz2_md, rf3, and protenix. Additional comparison presets include protenix_dual, rf3_protenix, and RF3 variants. Single-job presets default to gpu_count = 8, so on an 8-GPU pod they use the whole machine.

Presets live in experiments/*.toml in your local checkout and on the pod at /home/dev/workspace/experiments/*.toml. To modify an experiment, edit or copy a preset locally, let ACTL sync it, then run it by name or path:

cp experiments/rf3_partial.toml experiments/my_rf3.toml
# edit experiments/my_rf3.toml locally
run_experiments --preset my_rf3

For one-off changes, use --set instead of editing TOML:

run_experiments rf3 --set jobs.rf3.gpu_count=4
run_experiments rf3 --set jobs.rf3.args.gradient-weights="0.0 0.01 0.02"

Presets usually declare gpu_count = N, not fixed GPU IDs. The runner assigns visible GPUs automatically in job order, so the same preset works on different pod sizes and fails fast if the pod has fewer visible GPUs than requested. Use explicit gpus = "0,1" only when you need to pin a job to specific devices; the runner validates those IDs before launching jobs.

Set DATA_DIR and PROTEINS_CSV explicitly for each run so they are captured in the shell history and launch logs. Checkpoints default to /mnt/diffuse-shared/raw/checkpoints when those files exist, results go to /mnt/diffuse-shared/results/sampleworks/<pod>/<target>/, and MSA caches go to /mnt/diffuse-shared/cache/sampleworks/msa. Override with RESULTS_DIR, MSA_CACHE_DIR, or model-specific checkpoint variables before running.

The ACTL image contains baked pixi environments under /app/.pixi. If your synced branch changes pyproject.toml or pixi.lock, run_experiments stops with a clear error instead of mutating the baked environment. For dependency debugging only, opt into an on-pod pixi update with RUNTIME_PIXI=1 run_experiments ...; reproducible scientist runs should use a rebuilt pixi-with-checkpoints:sampleworks image instead.

Docker

TODO: Docker container documentation

Development

We use Pixi to manage development environments and dependencies. Each model has its own environment, e.g. boltz-dev, protenix-dev, rf3-dev. To install dev dependencies and run tests:

pixi install -e [model]-dev    # add pytest, ruff, ty
pixi run -e [model]-dev all-tests  # run tests
pixi run test-all            # run all tests across all environments

Prek hooks (various formatting, ruff + ty type checking):

pixi run -e [model]-dev prek install
pixi run -e [model]-dev prek install --hook-type commit-msg
pixi run -e [model]-dev prek run --all-files

See tests/README.md for full testing instructions.

macOS (experimental)

To develop on OS X, ensure you have homebrew installed and run the following commands to install dependencies:

1. Install hatch and uv

   brew install hatch uv

2. Move/copy pyproject-hatch.toml to pyproject.toml
3. Use uvx hatch run <command> to run commands. Note the use of uvx instead of uv
4. Use uvx hatch run <env>:<command> to run commands in a specific environment <env>.

There are different (and as yet untested) environments for boltz. protenix won't currently work on a Mac due to the strict requirement of triton which requires an NVIDIA GPU. You may find similar issues with other environments. Debug as needed.

Commit Messages

This project uses Conventional Commits to automate versioning and changelog generation. Format:

<type>(<scope>): <summary>

Common types: feat, fix, docs, refactor, chore, test, perf. A commitizen pre-commit hook validates messages at commit time. See AGENTS.md for full details.