Ligand-Conditioned Potts Models

Note: this README was LLM-written and manually reviewed.

This repository contains ligand-conditioned Potts/model development code around AlphaFold3-new, AtomWorks, and allatom_design. The active environment is named Elix.

Repository Layout

• allatom_design/: main Python package for data, model, evaluation, and utility code.
• alphafold3/: editable AlphaFold3-new checkout used by Elix.
• atomworks/: editable AtomWorks checkout used by Elix.
• requirements_split/: uv/pip requirement splits for local and Sherlock installs.
• scripts/: local and Sherlock setup, training, evaluation, and preprocessing scripts.
• tests/: lightweight pytest coverage, including Elix install smoke tests.

Sherlock Setup

Use the Elix-specific Sherlock guide:

scripts/sherlock_scripts/jinho/setup/README_elix_sherlock.md

That guide covers the full flow:

1. Build elix.sif locally with Apptainer.
2. Copy the SIF to Sherlock.
3. Install the Sherlock uv environment under scratch.
4. Validate imports, GPU visibility, patched HMMER --seq_limit, and AF3 data smoke tests.
5. Submit jobs through wrap_sbatch_in_container_elix.sh.

The previous Sherlock setup is intentionally not documented here because it used a different AlphaFold3 dependency set.

AF3 Databases On Sherlock

AlphaFold3 requires a local genetic database directory. During Sherlock setup, fetch the AF3 databases once into scratch:

export REPO="${REPO:-$HOME/code/elix}"
export AF3_DB_DIR="${AF3_DB_DIR:-$SCRATCH/af3_databases}"

cd "$REPO/alphafold3"
bash fetch_databases.sh "$AF3_DB_DIR"

Run this in tmux, screen, or an sbatch job. The full database download and decompression is large and can take a long time.

After the fetch completes, verify the expected AF3 v3 database layout:

for f in \
  bfd-first_non_consensus_sequences.fasta \
  mgy_clusters_2022_05.fa \
  uniref90_2022_05.fa \
  uniprot_all_2021_04.fa \
  pdb_seqres_2022_09_28.fasta \
  rnacentral_active_seq_id_90_cov_80_linclust.fasta \
  nt_rna_2023_02_23_clust_seq_id_90_cov_80_rep_seq.fasta \
  rfam_14_9_clust_seq_id_90_cov_80_rep_seq.fasta; do
  test -s "$AF3_DB_DIR/$f" || echo "MISSING $f"
done
test -d "$AF3_DB_DIR/mmcif_files" || echo "MISSING mmcif_files"

If AF3 self-consistency fails with a missing ${DB_DIR}/... database path, rerun the fetch command above or verify that struct_pred_cfg.af3.inference_config.base.db_dir points to $AF3_DB_DIR. Before rerunning a failed sample, remove any partial AF3 prediction directory:

rm -rf "$RUN_DIR"/step_*/af3_ss_preds/"$DESIGNED_SAMPLE_ID"

Local Desktop Setup

The local desktop environment uses a uv-managed Python 3.12 virtual environment for the AlphaFold3-new tree. Micromamba is used only to provide zlib headers and libraries while building AlphaFold3. OpenStructure is not installed.

1. Paths And Helper Tools

export MODEL_DEV=/path/to/model-dev
export REPO="$MODEL_DEV/allatom-design"
export VENV="$MODEL_DEV/envs/uv/elix_local"
export UV_CACHE_DIR="$MODEL_DEV/cache/uv"
export BUILD_PREFIX="$MODEL_DEV/envs/micromamba/envs/af3_build_deps"
export MAMBA_ROOT_PREFIX="$MODEL_DEV/envs/micromamba"

# Use the uv binary on PATH, or set this to a known uv executable.
export UV_BIN=uv

Install zlib into the small build prefix:

micromamba create -y -p "$BUILD_PREFIX" -c conda-forge zlib

2. Create The uv Environment

mkdir -p "$(dirname "$VENV")" "$UV_CACHE_DIR"
"$UV_BIN" python install 3.12
"$UV_BIN" venv "$VENV" -p 3.12
"$UV_BIN" pip install --python "$VENV/bin/python" --upgrade pip setuptools wheel

3. Install PyTorch/PyG And Build Tools

cd "$REPO"

"$UV_BIN" pip install --python "$VENV/bin/python" \
  scikit_build_core pybind11 'cmake>=3.28' ninja setuptools_scm

"$UV_BIN" pip install --python "$VENV/bin/python" \
  -r requirements_split/local/uv-compatible-torch.txt \
  --extra-index-url https://download.pytorch.org/whl/cu126 \
  --index-strategy unsafe-best-match

"$UV_BIN" pip install --python "$VENV/bin/python" \
  -r requirements_split/local/pip-only-torch.txt \
  --index-strategy unsafe-best-match

4. Install AlphaFold3

cd "$REPO"

CMAKE_PREFIX_PATH="$BUILD_PREFIX" ZLIB_ROOT="$BUILD_PREFIX" \
  "$UV_BIN" pip install --python "$VENV/bin/python" \
  -e ./alphafold3 --no-build-isolation --index-strategy unsafe-best-match

"$VENV/bin/build_data"

5. Install AtomWorks And allatom-design

AtomWorks currently declares rdkit<2025.9, while AlphaFold3-new requires rdkit==2025.09.4. Keep the AlphaFold3-new dependency set and install AtomWorks with --no-deps.

cd "$REPO"

"$UV_BIN" pip install --python "$VENV/bin/python" \
  mashumaro==3.14 nglview==3.1.4 hydra-core wandb \
  'pandas>=2.2,<2.4' seaborn matplotlib torchtyping einops biopython ihm modelcif \
  biotite==1.2.0 hydride 'py3Dmol>=2.2.1,<3' 'pymol-remote>=0.0.5' \
  pyarrow==17.0.0 'cython>=3,<4' 'cytoolz>=0.12.3,<1' 'typer>=0.12.5,<1' \
  'jaxtyping>=0.2.17,<1' 'beartype>=0.18.0,<1' pathspec pytest \
  --index-strategy unsafe-best-match

"$UV_BIN" pip install --python "$VENV/bin/python" openbabel-wheel==3.1.1.22
"$UV_BIN" pip install --python "$VENV/bin/python" -e ./atomworks --no-deps
"$UV_BIN" pip install --python "$VENV/bin/python" -e . --no-deps

JAX 0.9.1 needs a newer cuDNN runtime than the exact version declared by torch==2.7.0+cu126. Keep the newer cuDNN package in the environment:

"$UV_BIN" pip install --python "$VENV/bin/python" \
  nvidia-cudnn-cu12==9.22.0.52 --index-strategy unsafe-best-match

6. Use The Environment

source "$VENV/bin/activate"
cd "$REPO"

For local GPUs with compute capability 7.x, use XLA flash attention plus the XLA compatibility flag:

export XLA_FLAGS=--xla_disable_hlo_passes=custom-kernel-fusion-rewriter

python -m alphafold3.run_alphafold \
  --flash_attention_implementation=xla \
  ...

7. Validate

python -m compileall -q alphafold3/src/alphafold3 atomworks/src/atomworks allatom_design

python -c 'import jax, torch, rdkit, alphafold3, atomworks, allatom_design; print("jax", jax.__version__, jax.devices()); print("torch", torch.__version__, torch.cuda.is_available()); print("rdkit", rdkit.__version__); print("alphafold3", alphafold3.__file__); print("atomworks", atomworks.__file__)'

python -m pytest alphafold3/tests/test_runner_module.py -q

Evaluation Notes

AF3 template/input gap residues are filled from native residue names when those names are available in the evaluation sample metadata. Missing lookup entries still fall back to UNK.

AF3 v3.0.3 summary confidence JSON may not include interface_min_pae. The Elix metric code records that field as missing when absent; this warning alone does not mean AF3 prediction failed.

Expected pip check metadata conflicts:

• atomworks declares rdkit<2025.9, while AlphaFold3-new uses rdkit==2025.09.4.
• torch==2.7.0+cu126 declares nvidia-cudnn-cu12==9.5.1.17, while JAX 0.9.1 inference uses nvidia-cudnn-cu12==9.22.0.52.