diff --git a/AlphaFold3_of3.ipynb b/AlphaFold3_of3.ipynb index e20b5ffa4..2e87c3226 100644 --- a/AlphaFold3_of3.ipynb +++ b/AlphaFold3_of3.ipynb @@ -23,7 +23,7 @@ "\n", "MSA generation via the [ColabFold](https://github.com/sokrypton/ColabFold) MMseqs2 server — **no local databases required**.\n", "\n", - "Just fill in the **Input sequence(s)** cell and run all. Proteins, nucleic acids and ligands all go in a single box; attention/XLA flags are chosen automatically for your runtime (T4, newer GPUs, or CPU).\n", + "Just fill in the **Input sequences** cell and run all; attention/XLA flags are chosen automatically for your runtime (T4, newer GPUs, or CPU).\n", "\n", "**Citations:**\n", "- Abramson et al. (2024) AlphaFold 3. *Nature* [doi:10.1038/s41586-024-07487-w](https://doi.org/10.1038/s41586-024-07487-w)\n", @@ -44,27 +44,78 @@ "outputs": [], "source": [ "#@title Install dependencies (~3 mins)\n", - "%%time\n", - "import os, time, glob\n", - "\n", - "use_af3_weights = True #@param {type:\"boolean\"}\n", - "#@markdown - **use_af3_weights**: leave OFF to use the free, open **OpenFold3** weights. Turn ON to\n", - "#@markdown download the official **AlphaFold 3** weights directly - now public, no login or key required.\n", - "#@markdown Subject to the AF3 weights terms of use.\n", - "USE_NATIVE = use_af3_weights\n", + "import os, time, glob, shutil, sys, shlex\n", + "_T0 = time.time()\n", + "\n", + "model = \"openfold3\" #@param [\"openbind0\", \"openfold3\", \"alphafold3\"]\n", + "\n", + "persist_cache_to_drive = False #@param {type:\"boolean\"}\n", + "#@markdown - **persist_cache_to_drive**: keep the compiled model in Drive so the next\n", + "#@markdown session skips the recompile -- worth ~53 s (69 s cold vs 16 s warm on a\n", + "#@markdown 68-residue input). Never changes a result.\n", + "\n", + "# Set any form field from the environment, for runs outside Colab:\n", + "# AF3_NB_OVERRIDES='{\"model\": \"boltz2\"}'\n", + "import json as _json\n", + "for _k, _v in _json.loads(os.environ.get('AF3_NB_OVERRIDES', '{}')).items():\n", + " if _k in globals():\n", + " globals()[_k] = _v\n", + " print(f'override: {_k} = {_v!r}')\n", + "\n", + "# Bootstrap uv once; all subsequent package installation uses uv.\n", + "import subprocess\n", + "if shutil.which('uv') is None:\n", + " subprocess.run([sys.executable, '-m', 'pip', 'install', '-q', 'uv'], check=True)\n", + "_colabfold_source = os.environ.get('RGI_COLABFOLD_SOURCE',\n", + " 'git+https://github.com/sokrypton/ColabFold')\n", + "subprocess.run(['uv', 'pip', 'install', '--python', sys.executable, '--no-deps', _colabfold_source], check=True)\n", + "_toolkit_source = os.environ.get('RGI_TOOLKIT_SOURCE',\n", + " 'git+https://github.com/cddlab/rgi_toolkit.git@054cd4c81a7f198df12fdd72f88d5caa7f8d9615')\n", + "\n", + "VERSION = '3.1.11' # package and run_alphafold.py both come from this tag\n", "NATIVE_DIR = 'af3_native_weights'\n", "AF3_WEIGHTS_URL = 'https://storage.googleapis.com/alphafold3/af3.bin.zst'\n", + "AF2_DIR = 'af2_params'\n", + "IS_AF3 = (model == 'alphafold3')\n", + "IS_AF2 = model.startswith('af2_')\n", + "# int8 weights, expanded on load; AF2 and AF3 ship their own float32 files.\n", + "PRECISION = 'fp32' if (IS_AF3 or IS_AF2) else 'int8'\n", "\n", - "if not os.path.isfile('ALPHAFOLD3_READY'):\n", + "\n", + "def _sh(cmd, what):\n", + " \"\"\"Run a shell command, raising if it fails.\"\"\"\n", + " if cmd.startswith('uv pip install --system '):\n", + " cmd = cmd.replace('--system ', '--python ' + shlex.quote(sys.executable) + ' ', 1)\n", + " if os.system(cmd) != 0:\n", + " raise RuntimeError(f'{what} failed. The output is above.')\n", + "\n", + "\n", + "if not os.path.isfile(f'ALPHAFOLD3_READY_{VERSION}'):\n", " print('Installing packages...')\n", - " os.system(\"pip install -q 'jax[cuda12]==0.10.1' dm-haiku==0.0.17 rdkit==2025.9.4 \\\n", - " zstandard awscli tokamax==0.0.11 py3Dmol py2Dmol\")\n", - " os.system(\"pip install -q --no-deps 'https://github.com/sokrypton/alphafold3/releases/download/v3.1.5/alphafold3_open-3.1.5-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl'\")\n", - " base = 'https://raw.githubusercontent.com/sokrypton/alphafold3/refs/heads/main'\n", - " for script in ['run_alphafold', 'convert_of3_weights']:\n", - " os.system(f'wget -q -nc {base}/{script}.py')\n", - " os.system('touch ALPHAFOLD3_READY')\n", - " print('Packages installed.')\n", + " # Installed with --no-deps, so the package's own imports are listed here.\n", + " # Letting pip resolve them would re-download jax and the CUDA stack.\n", + " _sh(\"uv pip install --system -q dm-haiku==0.0.17 rdkit==2025.9.4 \"\n", + " \"tokamax==0.0.11 ml_collections\", 'installing dependencies')\n", + " _sh(\"uv pip install --system -q git+https://github.com/sokrypton/py2Dmol.git\", # wheel lags the repo\n", + " 'installing py2Dmol')\n", + " if IS_AF2:\n", + " os.system(\"apt-get -qq install -y aria2 > /dev/null 2>&1\") # AF2's tar is 5.3 GB\n", + " # Retried: PyPI's index can lag a just-published release by a few minutes.\n", + " for _try in range(4):\n", + " if os.system(f'uv pip install --python {shlex.quote(sys.executable)} -q --no-deps alphafold3-colabfold=={VERSION}') == 0:\n", + " break\n", + " print(f'pip could not find {VERSION} yet; retrying in 20 s')\n", + " time.sleep(20)\n", + " else:\n", + " raise RuntimeError(f'could not install alphafold3-colabfold=={VERSION}')\n", + " # run_alphafold.py is a top-level script, not part of the package.\n", + " _sh(f'wget -q -O run_alphafold.py https://raw.githubusercontent.com'\n", + " f'/sokrypton/alphafold3/v{VERSION}/run_alphafold.py', 'fetching run_alphafold.py')\n", + " # haiku 0.0.17 still calls the moved jax.core.DropVar.\n", + " os.system(\"sed -i 's/jax.core.DropVar/jax.extend.core.DropVar/g' /usr/local/lib/python*/dist-packages/haiku/_src/jaxpr_info.py\")\n", + " import alphafold3 # confirms the install before anything depends on it\n", + " open(f'ALPHAFOLD3_READY_{VERSION}', 'w').close()\n", + " print(f'Packages installed ({alphafold3.__file__}).')\n", "\n", "# Patch tokamax so Ada/consumer GPUs (L4, A10, RTX 30/40; cc 8.6/8.9) fall back to XLA\n", "# kernels. tokamax enables its Triton kernels for ALL cc>=8.0 GPUs, but those kernels\n", @@ -84,37 +135,67 @@ "except Exception as _e:\n", " print(f'(tokamax patch skipped: {_e})')\n", "\n", - "# Weights: official AlphaFold 3 (public download) or OpenFold3 (default, free) - both in background\n", - "if USE_NATIVE:\n", - " if not os.path.isfile('NATIVE_WEIGHTS_DONE'):\n", - " print('Downloading official AlphaFold 3 weights (public, no login required)...')\n", - " os.makedirs(NATIVE_DIR, exist_ok=True)\n", - " for _f in glob.glob(f'{NATIVE_DIR}/*'): # keep exactly one model file in the dir\n", - " os.remove(_f)\n", - " os.system(f'(wget -q -O {NATIVE_DIR}/af3.bin.zst \"{AF3_WEIGHTS_URL}\"; touch NATIVE_WEIGHTS_DONE) &')\n", - "elif not os.path.isfile('WEIGHTS_DONE'):\n", - " print('Downloading OpenFold3 weights...')\n", - " os.system('(aws s3 cp s3://openfold/staging/of3-p2-155k.pt . --no-sign-request; touch WEIGHTS_DONE) &')\n", - "\n", - "# Build AF3 data files (background, independent of weights)\n", - "if not os.path.isfile('DATA_DONE'):\n", - " print('Building AF3 data files...')\n", - " os.system('(build_data; touch DATA_DONE) &')\n", - "\n", - "# Wait for background jobs\n", - "for sentinel in (['NATIVE_WEIGHTS_DONE', 'DATA_DONE'] if USE_NATIVE else ['WEIGHTS_DONE', 'DATA_DONE']):\n", + "# Weights, fetched in the background by the same code the run uses.\n", + "STAMP = f'WEIGHTS_DONE_{model}_{PRECISION}'\n", + "if IS_AF3 and not os.path.isfile(STAMP):\n", + " # DeepMind's own release, subject to the AlphaFold 3 terms of use, which\n", + " # run_alphafold prints at startup. A copy you already have in NATIVE_DIR is\n", + " # used as-is.\n", + " os.makedirs(NATIVE_DIR, exist_ok=True)\n", + " if glob.glob(f'{NATIVE_DIR}/*.bin.zst'):\n", + " open(STAMP, 'w').close()\n", + " print(f'Using the AlphaFold 3 parameters already in {NATIVE_DIR}/.')\n", + " else:\n", + " print('Downloading AlphaFold 3 parameters (~1 GB)...')\n", + " os.system(f'(wget -O {NATIVE_DIR}/af3.bin.zst \"{AF3_WEIGHTS_URL}\"'\n", + " f' > {STAMP}.log 2>&1 && touch {STAMP}) &')\n", + "elif not (IS_AF3 or os.path.isfile(STAMP)):\n", + " _script, _args = ('prefetch_af2.py', AF2_DIR) if IS_AF2 else (\n", + " 'prefetch_weights.py', f'{model} {PRECISION}')\n", + " print(f'Downloading {\"official AlphaFold 2 parameters (CC BY 4.0)\" if IS_AF2 else model} weights...')\n", + " with open(_script, 'w') as fh:\n", + " fh.write('import sys\\n'\n", + " 'from alphafold3.model import weights\\n'\n", + " + ('print(weights.ensure_af2_params(sys.argv[1]))\\n' if IS_AF2 else\n", + " 'print(weights.ensure_weights(sys.argv[1], None, precision=sys.argv[2]))\\n'))\n", + " os.system(f'(python {_script} {_args} > {STAMP}.log 2>&1 && touch {STAMP}) &')\n", + "\n", + "# /tmp is wiped with the VM, so a fresh session recompiles (~53 s); Drive survives.\n", + "CACHE_DIR = os.path.join(__import__('tempfile').gettempdir(), 'af3_cache')\n", + "if persist_cache_to_drive:\n", + " try:\n", + " from google.colab import drive\n", + " drive.mount('/content/drive')\n", + " CACHE_DIR = '/content/drive/MyDrive/.af3_cache'\n", + " os.makedirs(CACHE_DIR, exist_ok=True)\n", + " print(f'Compile cache: {CACHE_DIR} (survives this session)')\n", + " except Exception as _e:\n", + " print(f'(Drive mount failed, using {CACHE_DIR}: {_e})')\n", + "\n", + "\n", + "def _await(sentinel, limit=1200):\n", + " \"\"\"Wait for a background job, reporting its log if it never finishes.\"\"\"\n", + " t0 = time.time()\n", " while not os.path.isfile(sentinel):\n", + " if time.time() - t0 > limit:\n", + " log = f'{sentinel}.log'\n", + " tail = open(log).read()[-1500:] if os.path.isfile(log) else '(no output captured)'\n", + " raise RuntimeError(f'{sentinel} did not appear within {limit} s. '\n", + " f'Tail of {log}:\\n{tail}')\n", " time.sleep(5)\n", - " print(f'{sentinel} ✓')\n", + " print(f'{sentinel} \\u2713 ({time.time() - t0:.0f} s)')\n", "\n", - "if USE_NATIVE and os.path.getsize(f'{NATIVE_DIR}/af3.bin.zst') < 1_000_000:\n", - " raise RuntimeError('AlphaFold 3 weights download failed or incomplete - re-run this cell.')\n", "\n", - "if not USE_NATIVE and not os.path.isfile('af3_converted_weights/of3_ported_weights.bin.zst'):\n", - " print(\"converting openfold3 weights...\")\n", - " os.system(\"python convert_of3_weights.py --of3_checkpoint of3-p2-155k.pt --output_dir af3_converted_weights\")\n", + "_await(STAMP)\n", "\n", - "print('Setup complete! ' + ('Using official AlphaFold 3 weights.' if USE_NATIVE else 'Using OpenFold3 weights.'))" + "if IS_AF3 and os.path.getsize(f'{NATIVE_DIR}/af3.bin.zst') < 1_000_000:\n", + " raise RuntimeError('the AlphaFold 3 download is incomplete - re-run this cell.')\n", + "\n", + "print(f'Setup complete! Model: {model}.')\n", + "if model == 'chai1':\n", + " print('NOTE: chai-1 is running WITHOUT ESM2 embeddings, which are most of its token\\n'\n", + " ' features. Expect worse structures than chai-lab itself produces.')\n", + "print(f'Setup took {time.time() - _T0:.0f} s.')\n" ] }, { @@ -122,8 +203,8 @@ "execution_count": null, "id": "input", "metadata": { - "id": "input", - "cellView": "form" + "cellView": "form", + "id": "input" }, "outputs": [], "source": [ @@ -147,6 +228,13 @@ "#@markdown - `seeds`: comma-separated, e.g. `1,2,3`.\n", "#@markdown - `on_existing`: `overwrite` replaces this job's previous results; `skip` keeps them.\n", "\n", + "# Headless overrides -- see the install cell.\n", + "import json as _json, os as _os\n", + "for _k, _v in _json.loads(_os.environ.get('AF3_NB_OVERRIDES', '{}')).items():\n", + " if _k in globals():\n", + " globals()[_k] = _v\n", + " print(f'override: {_k} = {_v!r}')\n", + "\n", "# Split a box into entries: collapse colon runs, drop whitespace, skip empties\n", "def split_entries(s):\n", " s = re.sub(r':+', ':', s).strip(':')\n", @@ -158,6 +246,24 @@ "ccd_codes = [e.upper() for e in split_entries(ligand_ccd)]\n", "smiles_strs = split_entries(ligand_smiles) # case-sensitive: leave as typed\n", "\n", + "# Fetch chemical definitions for the components this input names, from\n", + "# files.rcsb.org (~0.6 s). A code that is not fetched raises when folding.\n", + "with open('prefetch_ccd.py', 'w') as fh:\n", + " fh.write('import sys, os, importlib.metadata as md\\n'\n", + " 'from alphafold3.constants import ccd_fetch\\n'\n", + " 'root = os.path.dirname(md.distribution(\"alphafold3-colabfold\")'\n", + " '.locate_file(\"alphafold3\"))\\n'\n", + " 'conv = os.path.join(root, \"alphafold3\", \"constants\", \"converters\")\\n'\n", + " 'os.makedirs(conv, exist_ok=True)\\n'\n", + " 'ccd_fetch.write_pickles(ccd_fetch.codes_for_input(extra=sys.argv[1:]),\\n'\n", + " ' os.path.join(conv, \"ccd.pickle\"),\\n'\n", + " ' os.path.join(conv, \"chemical_component_sets.pickle\"),\\n'\n", + " ' libcifpp_dir=os.path.join(root, \"share\", \"libcifpp\"))\\n')\n", + "print(f'Fetching the CCD: 35 standard residues'\n", + " + (f' + {\", \".join(ccd_codes)}' if ccd_codes else '') + ' ...')\n", + "if os.system('python prefetch_ccd.py ' + ' '.join(ccd_codes)) != 0:\n", + " raise RuntimeError('could not build the CCD tables; see the output above')\n", + "\n", "# Build AF3 chain entities (IDs A, B, C, ... in canonical order)\n", "CHAIN_IDS = list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz')\n", "chains, prot_groups, idx = [], {}, 0\n", @@ -214,7 +320,7 @@ "jobname = basejob + '_' + hashlib.sha1(flat.encode()).hexdigest()[:5]\n", "\n", "# Input JSON goes to a temp dir; ALL results land in ONE folder: af3_output//\n", - "INPUT_DIR = '/tmp/af3_inputs'\n", + "INPUT_DIR = os.path.join(__import__('tempfile').gettempdir(), 'af3_inputs')\n", "OUTPUT_DIR = 'af3_output'\n", "job_dir = f'{OUTPUT_DIR}/{jobname}'\n", "\n", @@ -231,32 +337,120 @@ " json.dump(fold_input, f, indent=2)\n", "\n", "print(f'Job \"{jobname}\" -> results will be written to {job_dir}/')\n", - "fold_input\n" + "fold_input\n", + "\n", + "from colabfold.rgi.notebook import chain_table\n", + "print('Chain IDs for restraints (residue numbering starts at 1 in each chain):')\n", + "for row in chain_table(fold_input):\n", + " print(row)\n" ] }, { "cell_type": "code", "execution_count": null, - "id": "run", "metadata": { - "id": "run", + "id": "rgi-restraints", "cellView": "form" }, "outputs": [], + "source": [ + "#@title RGI: Making restraint-guided protein folding inference accessible to all! (optional)\n", + "#@markdown Fill in this form before **Runtime → Run all**. Leave `use_rgi` off for vanilla. Leave unused restraint\n", + "#@markdown fields empty.\n", + "use_rgi = False #@param {type:\"boolean\"}\n", + "#@markdown **distance** — enter two atom selections and their target distance in Å. Example: `chain A and resid\n", + "#@markdown 1 to 10` and `chain A and resid 40 to 50`, with `target_distance = 25`.\n", + "distance_atom_selection1 = \"\" #@param {type:\"string\"}\n", + "distance_atom_selection2 = \"\" #@param {type:\"string\"}\n", + "target_distance = 25 #@param {type:\"raw\"}\n", + "#@markdown **conformer** — enter chain IDs such as `B,C`, or `ligands` for all ligand chains. Leave empty to\n", + "#@markdown disable.\n", + "conformer_chains = \"\" #@param {type:\"string\"}\n", + "#@markdown **angle** — enter three atom selections and the target in degrees. Selection 2 is the vertex. Leave\n", + "#@markdown all three empty to disable.\n", + "angle_atom_selection1 = \"\" #@param {type:\"string\"}\n", + "angle_atom_selection2 = \"\" #@param {type:\"string\"}\n", + "angle_atom_selection3 = \"\" #@param {type:\"string\"}\n", + "target_angle = 90 #@param {type:\"raw\"}\n", + "#@markdown **custom** — set named selections and an energy expression, e.g. `harmonic(distance(A, B), 25.0)`.\n", + "#@markdown Leave the energy empty to disable.\n", + "custom_selections = {\"A\": \"chain A and resid 1 to 10\", \"B\": \"chain A and resid 40 to 50\"} #@param {type:\"raw\"}\n", + "custom_energy = \"\" #@param {type:\"string\"}\n", + "#@markdown **RMSD** — upload a reference using Colab's **Files → Upload** and enter its filename in either\n", + "#@markdown `ref_pdb` or `ref_cif`. Leave both filenames empty to disable. Empty target/ref selections mean the\n", + "#@markdown whole structure; e.g. `chain A and name CA` selects C-alpha atoms.\n", + "ref_pdb = \"\" #@param {type:\"string\"}\n", + "ref_cif = \"\" #@param {type:\"string\"}\n", + "atom_selection_target = \"\" #@param {type:\"string\"}\n", + "atom_selection_ref = \"\" #@param {type:\"string\"}\n", + "target_rmsd = 0 #@param {type:\"raw\"}\n", + "#@markdown **Multiple restraints:** enter lists in the same fields, e.g. `distance_atom_selection2 = [\"chain A\n", + "#@markdown and resid 20\", \"chain A and resid 40\"]` and `target_distance = [15, 25]`. A single selection/value is\n", + "#@markdown shared across the entries. All non-single lists must have the same length.\n", + "#@markdown **Selections:** all fields accept the full RGI-toolkit selection language. Residues start at 1 within\n", + "#@markdown each chain. The input cell prints the chain IDs. Distance and angle use atom-group centroids.\n", + "#@markdown **Advanced (optional)** — an ordinary RGI-toolkit dict for additional restraints, flat-bottomed\n", + "#@markdown penalties, weights, activation windows or conformer settings. Lists append to the fields above. Use\n", + "#@markdown `{\"config_path\": \"restraints.yaml\"}` to read a file. [Examples and native config\n", + "#@markdown reference](https://github.com/cddlab/rgi_toolkit/blob/main/docs/colabfold.md).\n", + "restraints_config = {} #@param {type:\"raw\"}\n", + "#@markdown After changing this form, rerun this cell and the prediction cell, or use **Run all**. Settings are\n", + "#@markdown saved in the notebook itself.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "run", + "metadata": { + "cellView": "form", + "id": "run" + }, + "outputs": [], "source": [ "#@title Run AlphaFold 3\n", - "%%time\n", - "import os, shutil, subprocess, glob\n", + "import os, shutil, subprocess, glob, time\n", + "_T0 = time.time()\n", "\n", - "#@markdown Inference settings (defaults match AlphaFold 3 - increase only if needed):\n", + "#@markdown Defaults match AlphaFold 3; raise only if needed.\n", "num_recycles = 10 #@param {type:\"integer\"}\n", "num_diffusion_samples = 5 #@param {type:\"integer\"}\n", - "#@markdown - `num_recycles`: refinement passes through the network (default 10). More can help large/hard targets, but is slower.\n", - "#@markdown - `num_diffusion_samples`: candidate structures generated per seed (default 5). Total models = seeds x samples.\n", + "#@markdown - `num_recycles`: refinement passes; more helps hard targets, costs time.\n", + "#@markdown - `num_diffusion_samples`: structures per seed, so total = seeds x samples.\n", + "\n", + "# Headless overrides -- see the install cell.\n", + "import json as _json, os as _os\n", + "for _k, _v in _json.loads(_os.environ.get('AF3_NB_OVERRIDES', '{}')).items():\n", + " if _k in globals():\n", + " globals()[_k] = _v\n", + " print(f'override: {_k} = {_v!r}')\n", "\n", "num_recycles = max(1, int(num_recycles))\n", "num_diffusion_samples = max(1, int(num_diffusion_samples))\n", "\n", + "# Read the standard Colab fields before starting inference.\n", + "if globals().get('use_rgi', False):\n", + " from colabfold.rgi import require_supported_model\n", + " require_supported_model(model)\n", + " subprocess.run(['uv', 'pip', 'install', '--python', sys.executable, _toolkit_source, 'jaxopt'], check=True)\n", + " from rgi_toolkit.notebook_colab import config_from_fields\n", + " rgi_config, rgi_conformer_chains = config_from_fields(globals())\n", + "else:\n", + " use_rgi = False\n", + " rgi_config, rgi_conformer_chains = None, ''\n", + "\n", + "# Rebuild the input even when this cell is rerun after switching RGI off.\n", + "from colabfold.rgi.notebook import name_job, prepare_input\n", + "fold_input = prepare_input(fold_input, model, use_rgi=use_rgi, config=rgi_config,\n", + " conformer_chains=rgi_conformer_chains)\n", + "jobname = name_job(fold_input, model, base=basejob, num_recycles=num_recycles,\n", + " num_diffusion_samples=num_diffusion_samples)\n", + "fold_input['name'] = jobname\n", + "job_dir = os.path.join(OUTPUT_DIR, jobname)\n", + "json_path = os.path.join(INPUT_DIR, f'{jobname}.json')\n", + "with open(json_path, 'w') as _f:\n", + " json.dump(fold_input, _f, indent=2)\n", + "\n", "os.makedirs(OUTPUT_DIR, exist_ok=True)\n", "\n", "# Re-run policy (one folder per job, no timestamped duplicates):\n", @@ -267,10 +461,7 @@ "if run_it:\n", " shutil.rmtree(job_dir, ignore_errors=True) # start clean so exactly one folder is produced\n", "\n", - "# Pick attention impl + XLA flags from the actual device.\n", - "# Triton/cuDNN flash attention need Ampere (compute capability >= 8.0);\n", - "# 7.x GPUs (T4=7.5, V100=7.0) and CPU use the portable XLA path, and 7.x\n", - "# additionally needs the XLA flag that disables the custom-kernel fusion pass.\n", + "# Attention implementation and XLA flags, chosen from the device.\n", "def detect_device():\n", " try:\n", " out = subprocess.run(\n", @@ -321,21 +512,16 @@ "\n", "print('XLA_FLAGS =', os.environ.get('XLA_FLAGS', '(unset)'))\n", "\n", - "# Weights: use the official AlphaFold 3 weights if they were fetched, else OpenFold3.\n", - "if os.path.isdir('af3_native_weights') and glob.glob('af3_native_weights/*.bin*'):\n", - " model_dir, use_of3 = 'af3_native_weights', False\n", - " print('Weights: official AlphaFold 3')\n", - "else:\n", - " model_dir, use_of3 = 'af3_converted_weights', True\n", - " print('Weights: OpenFold3')\n", + "# Ported models find their own cache, so --model_dir is only for AF2 and AF3.\n", + "print(f'Model: {model}')\n", "\n", "cmd = [\n", - " 'python', 'run_alphafold.py',\n", + " sys.executable, *(['-m', 'colabfold.rgi'] if use_rgi else []), 'run_alphafold.py',\n", " f'--json_path={json_path}',\n", - " f'--model_dir={model_dir}',\n", + " f'--model={model}',\n", " '--norun_data_pipeline',\n", " f'--output_dir={OUTPUT_DIR}',\n", - " '--cache_dir=/tmp/af3_cache',\n", + " f'--cache_dir={CACHE_DIR}',\n", " '--force_output_dir', # reuse af3_output// instead of a timestamped copy\n", " f'--flash_attention_implementation={flash_impl}',\n", " f'--num_recycles={num_recycles}',\n", @@ -343,18 +529,47 @@ "]\n", "if msa_mode == 'mmseqs2_server':\n", " cmd.append('--use_msa_server')\n", + "# chai-1 and ESMFold2 fold from a language model, downloaded on first use.\n", + "# Without it they are a different model, not a slightly worse one.\n", + "if model == 'chai1' or model.startswith('esmfold2'):\n", + " cmd.append('--use_esm_embeddings')\n", "if nojit:\n", " cmd.append('--nojit')\n", - "if use_of3:\n", - " cmd.append('--of3_weights')\n", + "if IS_AF3 or IS_AF2:\n", + " cmd.append(f'--model_dir={AF2_DIR if IS_AF2 else NATIVE_DIR}')\n", "\n", - "cmd = ' '.join(cmd)\n", + "import shlex\n", + "print('Command:', shlex.join(cmd))\n", "if run_it:\n", " print(cmd)\n", - " !{cmd}\n", - " print(f'\\nDone -> {job_dir}/')\n", + " # Popen rather than `!`: streams the output and gives an exit status.\n", + " _p = subprocess.Popen(cmd, stdout=subprocess.PIPE,\n", + " stderr=subprocess.STDOUT, text=True, bufsize=1)\n", + " for _line in _p.stdout:\n", + " print(_line, end='')\n", + " _rc = _p.wait()\n", + " _cifs = glob.glob(f'{job_dir}/**/*.cif', recursive=True)\n", + " if _rc != 0 or not _cifs:\n", + " raise RuntimeError(\n", + " f'the fold FAILED (exit {_rc}, {len(_cifs)} structures written). '\n", + " 'The output above is the whole story; scroll up for the error.')\n", + " print(f'\\nDone -> {job_dir}/ ({len(_cifs)} structures, '\n", + " f'{time.time() - _T0:.0f} s)')\n", "else:\n", - " print(f'Skipping: results already exist in {job_dir}/ (set on_existing=overwrite to recompute).')\n" + " print(f'Skipping: results already exist in {job_dir}/ (set on_existing=overwrite to recompute).')\n", + "\n", + "if use_rgi:\n", + " _report_path = os.path.join(job_dir, 'rgi_report.json')\n", + " if not os.path.isfile(_report_path):\n", + " raise RuntimeError('RGI report is missing; the run cannot be verified. Re-run with on_existing=overwrite.')\n", + " _rgi_report = json.load(open(_report_path))\n", + " from colabfold.citations import write_rgi_citation\n", + " write_rgi_citation(job_dir)\n", + " for _record in _rgi_report['seeds']:\n", + " print('RGI seed', _record['seed'], 'restraint counts:', _record['inventory'])\n", + " for _distance in _record['distances']:\n", + " print('Measured distances (Angstrom):', _distance['distances_angstrom'],\n", + " 'target/bounds:', _distance['lower'], _distance['upper'])\n" ] }, { @@ -419,7 +634,8 @@ " + (' (frames - press play)' if load_as_frames else ' (use the dropdown to switch)'))\n", "\n", "viewer = py2Dmol.view(size=(viewer_size, viewer_size),\n", - " pae=True, autoplay=load_as_frames)\n", + " pae=True, autoplay=load_as_frames,\n", + " style=\"cartoon\")\n", "for i, cif in enumerate(cifs, start=1):\n", " pae = load_pae(cif)\n", " if load_as_frames:\n", @@ -434,8 +650,8 @@ "execution_count": null, "id": "plots", "metadata": { - "id": "plots", - "cellView": "form" + "cellView": "form", + "id": "plots" }, "outputs": [], "source": [ @@ -551,26 +767,42 @@ "source": [ "# Instructions \n", "\n", - "**Quick start**\n", - "1. Fill in the sequence(s) in the **Input sequence(s)** cell.\n", - "2. Press **Runtime → Run all**.\n", - "3. The install cell (first run only) downloads OpenFold3 weights and builds AF3 data files in the background — subsequent runs reuse them.\n", + "**Quick start:** pick a **model** in the install cell, fill in the sequence(s), then\n", + "**Runtime → Run all**. The first run downloads that model's weights; later runs reuse\n", + "them, and each model has its own cache so switching back is instant.\n", "\n", "---\n", "\n", - "## Model weights\n", - "\n", - "By default the notebook uses the open **OpenFold3** weights (Apache-2.0, no restrictions). To run the **official AlphaFold 3** weights instead, tick the **use_af3_weights** box in the install cell. The weights are now publicly downloadable (no login, key, or DeepMind approval required) and are fetched directly into `af3_native_weights/`, then used automatically (the run cell detects them and drops the `--of3_weights` flag). Leaving **use_af3_weights** off keeps the OpenFold3 path.\n", - "\n", - "Official AlphaFold 3 parameters are subject to the [AF3 weights terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md); run_alphafold prints a reminder of these terms when they are used.\n", - "\n", + "## Models\n", + "\n", + "Ported weights come from [sokrypton/af3-any-model](https://huggingface.co/sokrypton/af3-any-model).\n", + "\"Tower\" is a protein language model fetched separately on first use.\n", + "\n", + "**RGI**: **Yes** supports `use_rgi = True`; **No** requires `use_rgi = False`.\n", + "\n", + "| model | weights | licence | RGI | notes |\n", + "|---|---|---|---|---|\n", + "| `openbind0` | [OpenFold3 v0.5.0 \"OpenBind\"](https://github.com/aqlaboratory/openfold-3/releases/tag/v0.5.0) (AlQuraishi Lab) | Apache-2.0 | Yes | The current OpenFold3 release, with optional RGI. |\n", + "| `openfold3` | [OpenFold3 preview-2](https://github.com/aqlaboratory/openfold) (AlQuraishi Lab) | Apache-2.0 | Yes | The earlier preview, kept because earlier results used it. |\n", + "| `boltz2` | [Boltz-2](https://github.com/jwohlwend/boltz) (Wohlwend et al.) | MIT | Yes | Strong on ligands; keeps a modified residue as one token. |\n", + "| `protenix2` | [Protenix-v2](https://github.com/bytedance/Protenix) (ByteDance) | Apache-2.0 | Yes | The widest trunk here (pair 256), so the slowest. |\n", + "| `rosettafold3` | [RoseTTAFold3](https://github.com/RosettaCommons/foundry) (RosettaCommons) | BSD-3-Clause | Yes | Carries chirality features; handles D-amino acids. |\n", + "| `chai1` | [chai-1](https://github.com/chaidiscovery/chai-lab) (Chai Discovery) | Apache-2.0 | Yes | Folds from ESM2 3B, fetched and run automatically. |\n", + "| `intellifold2` | [IntelliFold-v2](https://huggingface.co/intelligenAI/intellifold) (IntelliGen-AI) | Apache-2.0 | No | Widened channels (pair 512), largest ported download. |\n", + "| `opendde` | [OpenDDE](https://huggingface.co/aurekaresearch/OpenDDE) (Aureka Research) | Apache-2.0 | Yes | Runs its diffusion on an expanded structural-token set. |\n", + "| `esmfold2` | [ESMFold2](https://huggingface.co/biohub/ESMFold2) (Chan Zuckerberg Biohub) | MIT | Yes | Folds from ESM-C instead of an MSA — single sequence, no search. |\n", + "| `esmfold2_lm600m` | ESMFold2, 600M tower | MIT | Yes | No confidence head. |\n", + "| `esmfold2_lm300m` | ESMFold2, 300M tower | MIT | Yes | No confidence head. |\n", + "| `af2_ptm` | AlphaFold 2 monomer pTM (DeepMind) | CC BY 4.0 | No | **Protein only** — a ligand or nucleotide raises rather than quietly folding the rest. Templates use the model_1/model_2 parameter sets. |\n", + "| `af2_multimer` | AlphaFold 2 multimer v3 (DeepMind) | CC BY 4.0 | No | Protein only, as above. |\n", + "| `alphafold3` | Google DeepMind's own parameters | [AF3 terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md) | Yes | DeepMind's public release, downloaded on first use (~1 GB). Its terms govern the weights and the outputs; run_alphafold prints them at startup. |\n", "---\n", "\n", - "## Sequence input\n", + "## Input\n", "\n", - "Each molecule type has its own box. Within a box, separate multiple chains with `:`.\n", + "Each molecule type has its own box; within a box, separate chains with `:`.\n", "\n", - "| Box | What goes in it | Example |\n", + "| box | contents | example |\n", "|---|---|---|\n", "| **protein** | amino-acid sequence(s) | `MKTAY...` or `SEQ1:SEQ2` |\n", "| **dna** | DNA sequence(s) | `CGCGAATTCGCG` |\n", @@ -578,57 +810,57 @@ "| **ligand_ccd** | ligand(s) by PDB CCD code | `ATP:MG:HEM` |\n", "| **ligand_smiles** | ligand(s) by SMILES | `CC(=O)Oc1ccccc1C(=O)O` |\n", "\n", - "Chains are assigned IDs A, B, C, … following AlphaFold 3's canonical order (protein → RNA → DNA → ligand; CCD ligands before SMILES ligands). Mix freely across boxes to build a complex — e.g. a protein in **protein**, `AUGCAUGC` in **rna**, and `ATP` in **ligand_ccd**.\n", - "\n", - "- **Homo-oligomers**: identical protein sequences are merged automatically, so `SEQ:SEQ` = homodimer, `SEQ:SEQ:SEQ` = homotrimer.\n", - "- Protein / DNA / RNA sequences and CCD codes are upper-cased automatically; **SMILES are left exactly as typed** (case is meaningful in SMILES).\n", - "- Spaces and newlines inside an entry are ignored, and **extra colons are forgiven** — `SEQ1::::SEQ2` is the same as `SEQ1:SEQ2`. Leave a box empty if unused.\n", - "- *Note:* because `:` separates entries, an atom-mapped SMILES that itself contains a colon (e.g. `[C:1]`) isn't supported via the box — use a raw AF3 JSON for that edge case.\n", - "\n", - "## Seeds\n", - "\n", - "Enter one or more model seeds in the **seeds** box, comma-separated (e.g. `1,2,3`). Each seed is an independent prediction (more seeds = more sampling, more runtime). Non-numeric characters are ignored and duplicates are dropped, so `1, 1, foo, 7` becomes seeds `1` and `7`.\n", + "Mix boxes freely to build a complex. Chain IDs A, B, C… follow AlphaFold 3's canonical\n", + "order (protein → RNA → DNA → ligand). Identical protein sequences are merged, so\n", + "`SEQ:SEQ` is a homodimer. Sequences and CCD codes are upper-cased; **SMILES are left\n", + "as typed**. Whitespace and extra colons are forgiven (`SEQ1::::SEQ2` = `SEQ1:SEQ2`) —\n", + "which is also why an atom-mapped SMILES containing `:` needs a raw AF3 JSON instead.\n", "\n", - "## MSA modes\n", + "**seeds**: comma-separated, one prediction each (`1,2,3`). Junk and duplicates are\n", + "dropped. **msa_mode**: `mmseqs2_server` queries the public\n", + "[ColabFold](https://colabfold.mmseqs.com/) API (protein only — RNA/DNA always run\n", + "MSA-free); `single_sequence` skips it, faster and less accurate.\n", "\n", - "- **`mmseqs2_server`** *(recommended)*: queries the public [ColabFold](https://colabfold.mmseqs.com/) MMseqs2 API. Covers UniRef30 + environmental sequences for proteins. RNA/DNA chains always run MSA-free (ColabFold is protein-only).\n", - "- **`single_sequence`**: no MSA, query sequence only. Faster but less accurate, especially for monomers with close homologs.\n", + "## Output\n", "\n", - "## Output files (inside the downloaded zip)\n", - "\n", - "| File | Contents |\n", + "| file | contents |\n", "|---|---|\n", - "| `*.cif` | Best-ranked structure in mmCIF format. B-factor = pLDDT (0–100). |\n", - "| `*_confidences.json` | Per-residue pLDDT, PAE matrix, contact probs. |\n", + "| `*.cif` | Best-ranked structure. B-factor = pLDDT (0–100). |\n", + "| `*_confidences.json` | Per-residue pLDDT, PAE matrix, contact probabilities. |\n", "| `*_summary_confidences.json` | Mean pLDDT, pTM, ipTM, ranking score. |\n", - "| `*_ranking_scores.csv` | Ranking scores for all seed × sample combinations. |\n", - "| `seed-N_sample-M/` | Individual prediction directories (one per seed/sample). |\n", - "| `TERMS_OF_USE.md` | Apache 2.0 notice (OpenFold3 weights have no commercial restrictions). |\n", - "\n", - "## Interpreting confidence scores\n", + "| `*_ranking_scores.csv` | Every seed × sample combination. |\n", + "| `seed-N_sample-M/` | One directory per prediction. |\n", + "| `TERMS_OF_USE.md` | The licence for whichever weights you ran. |\n", "\n", - "- **pLDDT > 90**: very high confidence.\n", - "- **pLDDT 70–90**: confident, backbone generally reliable.\n", - "- **pLDDT 50–70**: low confidence, treat with caution.\n", - "- **pLDDT < 50**: very low, likely disordered or incorrect.\n", - "- **PAE**: lower values = confident relative positioning between residue pairs. Useful for assessing interface quality in complexes.\n", - "- **ipTM > 0.8**: strong evidence for a well-defined complex interface. **ipTM is `n/a` for single-chain jobs** (there is no interface to score).\n", + "pLDDT above 90 is very high, 70–90 reliable backbone, 50–70 doubtful, below 50 likely\n", + "disordered or wrong. Lower PAE means two residues are confidently placed *relative to\n", + "each other*, which is what to read for an interface. ipTM above 0.8 is a well-defined\n", + "complex interface, and is `n/a` for a single chain — there is no interface to score.\n", "\n", "## Troubleshooting\n", "\n", - "- **Check runtime type**: `Runtime → Change runtime type → GPU` (T4 is fine; A100/L4 are faster).\n", - "- **OOM error**: reduce sequence length or use a larger-memory GPU runtime.\n", - "- **MSA server timeout**: the public ColabFold server is rate-limited. Try again later or switch to `single_sequence` mode.\n", - "- **Download popup blocked**: disable your ad blocker.\n", - "- **Weight download slow**: the OF3 checkpoint is ~1.4 GB from AWS S3; the install cell waits automatically.\n", + "**OOM**: shorter sequence, or a larger GPU (`Runtime → Change runtime type`).\n", + "**MSA server timeout**: the public server is rate-limited — retry, or use\n", + "`single_sequence`. **Download popup blocked**: disable your ad blocker.\n", "\n", - "## License\n", + "## Licence\n", "\n", - "AlphaFold 3 source code and OpenFold3 weights are both licensed under [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0). Outputs generated with OpenFold3 weights are **not** subject to Google DeepMind's AlphaFold 3 Output Terms of Use and may be used freely, including commercially.\n", + "The AlphaFold 3 **source code** is [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0);\n", + "the **weights** are each their own, as listed in the model table. Outputs from the seven\n", + "Apache/MIT/BSD-licensed ported models are **not** subject to DeepMind's AlphaFold 3 Output\n", + "Terms of Use and may be used freely, including commercially. `alphafold3` is the exception:\n", + "its parameters and outputs carry DeepMind's own terms. Every run writes a\n", + "`TERMS_OF_USE.md` naming the licence that actually applies to it.\n", "\n", "## Bugs / feedback\n", "\n", - "Report issues at https://github.com/sokrypton/alphafold3/issues" + "https://github.com/sokrypton/alphafold3/issues\n", + "\n", + "## RGI: Making restraint-guided protein folding inference accessible to all! (optional)\n", + "\n", + "Before **Run all**, enable `use_rgi` and fill in the **[RGI form](#scrollTo=rgi-restraints)**.\n", + "The fields are visible before any cell runs. Leave unused selections empty.\n", + "See the [RGI-toolkit guide](https://github.com/cddlab/rgi_toolkit/blob/main/docs/colabfold.md) for examples and multiple restraints.\n" ] } ], @@ -640,13 +872,23 @@ "include_colab_link": true }, "kernelspec": { - "display_name": "Python 3", + "display_name": "Python 3 (ipykernel)", + "language": "python", "name": "python3" }, "language_info": { - "name": "python" + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" } }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/Boltz1.ipynb b/Boltz1.ipynb index b5389d7bf..d98487624 100644 --- a/Boltz1.ipynb +++ b/Boltz1.ipynb @@ -70,6 +70,11 @@ "#@markdown - Use `:` to specify multiple DNA sequences\n", "jobname = 'test' #@param {type:\"string\"}\n", "\n", + "msa_mode = \"mmseqs2_server\" #@param [\"mmseqs2_server\", \"single_sequence\"]\n", + "import json\n", + "for _k, _v in json.loads(os.environ.get('BOLTZ_NB_OVERRIDES', '{}')).items():\n", + " if _k in globals(): globals()[_k] = _v\n", + "\n", "# Clean up the query sequence and jobname\n", "query_sequence = \"\".join(query_sequence.split())\n", "ligand_input = \"\".join(ligand_input.split())\n", @@ -219,7 +224,11 @@ "# Optionally, print the output for verification\n", "#print(f\"Generated FASTA file '{queries_fasta}':\\n\")\n", "#for header, sequence in fasta_entries:\n", - "# print(f\"{header}\\n{sequence}\\n\")\n" + "# print(f\"{header}\\n{sequence}\\n\")\n", + "\n", + "print('Chain IDs for restraints:')\n", + "for header, sequence in fasta_entries:\n", + " print(header.split('|')[0].lstrip('>'), header.split('|')[1], len(sequence) if header.split('|')[1] == 'protein' else '')\n" ], "metadata": { "cellView": "form", @@ -228,6 +237,60 @@ "execution_count": null, "outputs": [] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "rgi-restraints", + "cellView": "form" + }, + "outputs": [], + "source": [ + "#@title RGI: Making restraint-guided protein folding inference accessible to all! (optional)\n", + "#@markdown Fill in this form before **Runtime → Run all**. Leave `use_rgi` off for vanilla. Leave unused restraint\n", + "#@markdown fields empty.\n", + "use_rgi = False #@param {type:\"boolean\"}\n", + "#@markdown **distance** — enter two atom selections and their target distance in Å. Example: `chain A and resid\n", + "#@markdown 1 to 10` and `chain A and resid 40 to 50`, with `target_distance = 25`.\n", + "distance_atom_selection1 = \"\" #@param {type:\"string\"}\n", + "distance_atom_selection2 = \"\" #@param {type:\"string\"}\n", + "target_distance = 25 #@param {type:\"raw\"}\n", + "#@markdown **conformer** — enter chain IDs such as `B,C`, or `ligands` for all ligand chains. Leave empty to\n", + "#@markdown disable.\n", + "conformer_chains = \"\" #@param {type:\"string\"}\n", + "#@markdown **angle** — enter three atom selections and the target in degrees. Selection 2 is the vertex. Leave\n", + "#@markdown all three empty to disable.\n", + "angle_atom_selection1 = \"\" #@param {type:\"string\"}\n", + "angle_atom_selection2 = \"\" #@param {type:\"string\"}\n", + "angle_atom_selection3 = \"\" #@param {type:\"string\"}\n", + "target_angle = 90 #@param {type:\"raw\"}\n", + "#@markdown **custom** — set named selections and an energy expression, e.g. `harmonic(distance(A, B), 25.0)`.\n", + "#@markdown Leave the energy empty to disable.\n", + "custom_selections = {\"A\": \"chain A and resid 1 to 10\", \"B\": \"chain A and resid 40 to 50\"} #@param {type:\"raw\"}\n", + "custom_energy = \"\" #@param {type:\"string\"}\n", + "#@markdown **RMSD** — upload a reference using Colab's **Files → Upload** and enter its filename in either\n", + "#@markdown `ref_pdb` or `ref_cif`. Leave both filenames empty to disable. Empty target/ref selections mean the\n", + "#@markdown whole structure; e.g. `chain A and name CA` selects C-alpha atoms.\n", + "ref_pdb = \"\" #@param {type:\"string\"}\n", + "ref_cif = \"\" #@param {type:\"string\"}\n", + "atom_selection_target = \"\" #@param {type:\"string\"}\n", + "atom_selection_ref = \"\" #@param {type:\"string\"}\n", + "target_rmsd = 0 #@param {type:\"raw\"}\n", + "#@markdown **Multiple restraints:** enter lists in the same fields, e.g. `distance_atom_selection2 = [\"chain A\n", + "#@markdown and resid 20\", \"chain A and resid 40\"]` and `target_distance = [15, 25]`. A single selection/value is\n", + "#@markdown shared across the entries. All non-single lists must have the same length.\n", + "#@markdown **Selections:** all fields accept the full RGI-toolkit selection language. Residues start at 1 within\n", + "#@markdown each chain. The input cell prints the chain IDs. Distance and angle use atom-group centroids.\n", + "#@markdown **Advanced (optional)** — an ordinary RGI-toolkit dict for additional restraints, flat-bottomed\n", + "#@markdown penalties, weights, activation windows or conformer settings. Lists append to the fields above. Use\n", + "#@markdown `{\"config_path\": \"restraints.yaml\"}` to read a file. [Examples and native config\n", + "#@markdown reference](https://github.com/cddlab/rgi_toolkit/blob/main/docs/colabfold.md).\n", + "restraints_config = {} #@param {type:\"raw\"}\n", + "#@markdown After changing this form, rerun this cell and the prediction cell, or use **Run all**. Settings are\n", + "#@markdown saved in the notebook itself.\n", + "#@markdown When switching `use_rgi`, rerun **Install dependencies** too to select the matching Boltz environment.\n" + ] + }, { "cell_type": "code", "execution_count": null, @@ -238,34 +301,39 @@ "outputs": [], "source": [ "#@title Install dependencies\n", - "%%time\n", - "import os\n", - "if not os.path.isfile(\"COLABFOLD_READY\"):\n", - " print(\"installing colabfold...\")\n", - " os.system(\"pip install -q --no-warn-conflicts 'colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold'\")\n", - " if os.environ.get('TPU_NAME', False) != False:\n", - " os.system(\"pip uninstall -y jax jaxlib\")\n", - " os.system(\"pip install --no-warn-conflicts --upgrade dm-haiku==0.0.10 'jax[cuda12_pip]'==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\")\n", - " os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabfold colabfold\")\n", - " os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/alphafold alphafold\")\n", - " # hack to fix TF crash\n", - " os.system(\"rm -f /usr/local/lib/python3.*/dist-packages/tensorflow/core/kernels/libtfkernel_sobol_op.so /usr/local/lib/python3.*/dist-packages/tensorflow/lite/python/*/*.so\")\n", - " os.system(\"touch COLABFOLD_READY\")\n", - "\n", - "if not os.path.isfile(\"BOLZ_READY\"):\n", - " os.system(\"apt-get install -y aria2\")\n", - " os.system(\"pip install -q --no-warn-conflicts boltz\")\n", - " os.system(\"mkdir weights\")\n", - " os.system(\"aria2c -d weights -x8 -s8 https://colabfold.steineggerlab.workers.dev/boltz1.ckpt\")\n", - " os.system(\"aria2c -d weights -x8 -s8 https://colabfold.steineggerlab.workers.dev/ccd.pkl\")\n", - " os.system(\"touch BOLZ_READY\")" + "import shutil, subprocess, sys\n", + "from pathlib import Path\n", + "use_rgi = globals().get('use_rgi', False)\n", + "if shutil.which('uv') is None:\n", + " subprocess.run([sys.executable, '-m', 'pip', 'install', '-q', 'uv'], check=True)\n", + "_colabfold_source = os.environ.get('RGI_COLABFOLD_SOURCE', 'git+https://github.com/sokrypton/ColabFold')\n", + "_toolkit_source = os.environ.get('RGI_TOOLKIT_SOURCE', 'git+https://github.com/cddlab/rgi_toolkit.git@054cd4c81a7f198df12fdd72f88d5caa7f8d9615')\n", + "subprocess.run(['uv', 'pip', 'install', '--python', sys.executable, _colabfold_source], check=True)\n", + "# Boltz requires Python 3.12; isolate it from Colab's Python and NumPy versions.\n", + "boltz_env = Path('.cache') / ('boltz-rgi' if use_rgi else 'boltz-vanilla')\n", + "if not (boltz_env / 'bin/python').exists():\n", + " subprocess.run(['uv', 'venv', '--python', '3.12', str(boltz_env)], check=True)\n", + "boltz_python = str(boltz_env / 'bin/python')\n", + "_boltz_source = ('git+https://github.com/cddlab/boltz_restr.git@rgi-integration' if use_rgi else 'boltz==2.2.1')\n", + "subprocess.run(['uv', 'pip', 'install', '--python', boltz_python, _boltz_source], check=True)\n", + "if use_rgi:\n", + " subprocess.run(['uv', 'pip', 'install', '--python', sys.executable, _toolkit_source], check=True)\n", + " subprocess.run(['uv', 'pip', 'install', '--python', boltz_python, _toolkit_source], check=True)\n" ] }, { "cell_type": "code", "source": [ "#@title Generate MSA with ColabFold\n", - "!colabfold_batch \"{queries_path}\" \"{jobname}\" --msa-only" + "if msa_mode == 'mmseqs2_server' and csv_entries:\n", + " from colabfold.colabfold import run_mmseqs2\n", + " for seq_id, seq in csv_entries:\n", + " msa_path = Path(jobname) / f'{seq_id}.a3m'\n", + " if not msa_path.exists():\n", + " a3m = run_mmseqs2([seq], str(Path(jobname) / (seq_id + '_msa')), use_env=True)[0]\n", + " msa_path.write_text(a3m)\n", + "else:\n", + " print('Single sequence: skipping the MSA server.')\n" ], "metadata": { "cellView": "form", @@ -278,7 +346,46 @@ "cell_type": "code", "source": [ "#@title Predict structure using boltz\n", - "!boltz predict --out_dir \"{jobname}\" \"{jobname}/{jobname}.fasta\" --cache weights" + "import yaml\n", + "num_recycles = 3 #@param {type:\"integer\"}\n", + "num_diffusion_samples = 1 #@param {type:\"integer\"}\n", + "seed = 42 #@param {type:\"integer\"}\n", + "# Read the standard Colab fields before starting inference.\n", + "if globals().get('use_rgi', False):\n", + " from rgi_toolkit.notebook_colab import config_from_fields\n", + " rgi_config, rgi_conformer_chains = config_from_fields(globals())\n", + "else:\n", + " use_rgi = False\n", + " rgi_config, rgi_conformer_chains = None, ''\n", + "\n", + "from colabfold.rgi.boltz import notebook_input\n", + "\n", + "# Use YAML so the predictor can read the shared RGI config and entity opt-ins.\n", + "boltz_input = notebook_input(fasta_entries, config=rgi_config if use_rgi else None,\n", + " conformer_chains=rgi_conformer_chains, single_sequence=(msa_mode == 'single_sequence'))\n", + "input_yaml = Path(jobname) / f'{jobname}.yaml'\n", + "input_yaml.write_text(yaml.safe_dump(boltz_input, sort_keys=False))\n", + "# Refresh RGI features without changing vanilla-only cache reuse.\n", + "boltz_output = Path(jobname) / f'boltz_results_{input_yaml.stem}'\n", + "rgi_cache_marker = boltz_output / '.rgi-enabled'\n", + "if use_rgi or rgi_cache_marker.exists():\n", + " import shutil\n", + " if boltz_output.exists():\n", + " shutil.rmtree(boltz_output)\n", + " if use_rgi:\n", + " boltz_output.mkdir(parents=True, exist_ok=True)\n", + " rgi_cache_marker.touch()\n", + "subprocess.run([boltz_python, '-m', 'boltz.main', 'predict', str(input_yaml),\n", + " '--model', 'boltz1', '--out_dir', jobname, '--cache', str(Path('.cache') / 'boltz-weights'),\n", + " '--recycling_steps', str(num_recycles), '--diffusion_samples', str(num_diffusion_samples),\n", + " '--seed', str(seed), '--override'], check=True)\n", + "import glob\n", + "if not glob.glob(f'{jobname}/boltz_results_{jobname}/predictions/{jobname}/*.cif'):\n", + " raise RuntimeError('Prediction produced no structure files. Check the log above.')\n", + "if use_rgi:\n", + " from colabfold.citations import write_rgi_citation\n", + " write_rgi_citation(boltz_output)\n", + "print('RGI ON: check nonzero built-spec counts and final energies above.' if use_rgi else 'Vanilla prediction complete.')\n" ], "metadata": { "cellView": "form", @@ -307,7 +414,8 @@ "with zipfile.ZipFile(zip_filename, 'w') as zipf:\n", " coverage_png_files = glob.glob(os.path.join(jobname, '*_coverage.png'))\n", " a3m_files = glob.glob(os.path.join(jobname, '*.a3m'))\n", - " for file in coverage_png_files + a3m_files:\n", + " bibtex_files = glob.glob(os.path.join(jobname, f'boltz_results_{jobname}', 'cite.bibtex'))\n", + " for file in coverage_png_files + a3m_files + [str(input_yaml)] + bibtex_files:\n", " arcname = os.path.basename(file) # Use only the file name\n", " zipf.write(file, arcname=arcname)\n", "\n", @@ -334,4 +442,4 @@ "outputs": [] } ] -} \ No newline at end of file +} diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb index 462a5b4b0..2bd387a5d 100644 --- a/ColabFold2_preview.ipynb +++ b/ColabFold2_preview.ipynb @@ -3,8 +3,8 @@ { "cell_type": "markdown", "metadata": { - "id": "view-in-github", - "colab_type": "text" + "colab_type": "text", + "id": "view-in-github" }, "source": [ "\"Open" @@ -20,7 +20,7 @@ "# ColabFold2 preview\n", "\n", "Predict protein, RNA, DNA and small-molecule structures with\n", - "[AlphaFold 3](https://www.nature.com/articles/s41586-024-07487-w) \u2014 and with thirteen\n", + "[AlphaFold 3](https://www.nature.com/articles/s41586-024-07487-w) — and with thirteen\n", "other sets of weights, all through one implementation. Pick a model in the install\n", "cell; the weights download themselves. MSAs come from the\n", "[ColabFold](https://github.com/sokrypton/ColabFold) MMseqs2 server, so no local\n", @@ -31,7 +31,8 @@ "\n", "**Please cite** AlphaFold 3 ([Abramson 2024](https://doi.org/10.1038/s41586-024-07487-w)),\n", "ColabFold ([Mirdita 2022](https://doi.org/10.1038/s41592-022-01488-1)), and whichever\n", - "model you ran.\n" + "model you ran. If you enable RGI, also cite RGI-Toolkit\n", + "([Hori 2026](https://doi.org/10.64898/2026.10.05.756905); [BibTeX](#rgi-citation)).\n" ] }, { @@ -45,7 +46,7 @@ "outputs": [], "source": [ "#@title Install dependencies (~35 s)\n", - "import os, time, glob, shutil, sys\n", + "import os, time, glob, shutil, sys, shlex\n", "_T0 = time.time()\n", "\n", "model = \"openbind0\" #@param [\"openbind0\", \"openfold3\", \"boltz2\", \"protenix2\", \"rosettafold3\", \"chai1\", \"intellifold2\", \"opendde\", \"esmfold2\", \"esmfold2_lm600m\", \"esmfold2_lm300m\", \"alphafold3\", \"af2_ptm\", \"af2_multimer\"]\n", @@ -63,7 +64,17 @@ " globals()[_k] = _v\n", " print(f'override: {_k} = {_v!r}')\n", "\n", - "VERSION = '3.1.11' # package and run_alphafold.py both come from this tag\n", + "# Bootstrap uv once; all subsequent package installation uses uv.\n", + "import subprocess\n", + "if shutil.which('uv') is None:\n", + " subprocess.run([sys.executable, '-m', 'pip', 'install', '-q', 'uv'], check=True)\n", + "_colabfold_source = os.environ.get('RGI_COLABFOLD_SOURCE',\n", + " 'git+https://github.com/sokrypton/ColabFold')\n", + "subprocess.run(['uv', 'pip', 'install', '--python', sys.executable, '--no-deps', _colabfold_source], check=True)\n", + "_toolkit_source = os.environ.get('RGI_TOOLKIT_SOURCE',\n", + " 'git+https://github.com/cddlab/rgi_toolkit.git@054cd4c81a7f198df12fdd72f88d5caa7f8d9615')\n", + "\n", + "VERSION = '3.1.11' # package and run_alphafold.py both come from this tag\n", "NATIVE_DIR = 'af3_native_weights'\n", "AF3_WEIGHTS_URL = 'https://storage.googleapis.com/alphafold3/af3.bin.zst'\n", "AF2_DIR = 'af2_params'\n", @@ -75,23 +86,25 @@ "\n", "def _sh(cmd, what):\n", " \"\"\"Run a shell command, raising if it fails.\"\"\"\n", + " if cmd.startswith('uv pip install --system '):\n", + " cmd = cmd.replace('--system ', '--python ' + shlex.quote(sys.executable) + ' ', 1)\n", " if os.system(cmd) != 0:\n", " raise RuntimeError(f'{what} failed. The output is above.')\n", "\n", "\n", - "if not os.path.isfile('ALPHAFOLD3_READY'):\n", + "if not os.path.isfile(f'ALPHAFOLD3_READY_{VERSION}'):\n", " print('Installing packages...')\n", " # Installed with --no-deps, so the package's own imports are listed here.\n", " # Letting pip resolve them would re-download jax and the CUDA stack.\n", - " _sh(\"pip install -q dm-haiku==0.0.17 rdkit==2025.9.4 \"\n", + " _sh(\"uv pip install --system -q dm-haiku==0.0.17 rdkit==2025.9.4 \"\n", " \"tokamax==0.0.11 ml_collections\", 'installing dependencies')\n", - " _sh(\"pip install -q git+https://github.com/sokrypton/py2Dmol.git\", # wheel lags the repo\n", + " _sh(\"uv pip install --system -q git+https://github.com/sokrypton/py2Dmol.git\", # wheel lags the repo\n", " 'installing py2Dmol')\n", " if IS_AF2:\n", " os.system(\"apt-get -qq install -y aria2 > /dev/null 2>&1\") # AF2's tar is 5.3 GB\n", " # Retried: PyPI's index can lag a just-published release by a few minutes.\n", " for _try in range(4):\n", - " if os.system(f'pip install -q --no-deps alphafold3-colabfold=={VERSION}') == 0:\n", + " if os.system(f'uv pip install --python {shlex.quote(sys.executable)} -q --no-deps alphafold3-colabfold=={VERSION}') == 0:\n", " break\n", " print(f'pip could not find {VERSION} yet; retrying in 20 s')\n", " time.sleep(20)\n", @@ -103,7 +116,7 @@ " # haiku 0.0.17 still calls the moved jax.core.DropVar.\n", " os.system(\"sed -i 's/jax.core.DropVar/jax.extend.core.DropVar/g' /usr/local/lib/python*/dist-packages/haiku/_src/jaxpr_info.py\")\n", " import alphafold3 # confirms the install before anything depends on it\n", - " os.system('touch ALPHAFOLD3_READY')\n", + " open(f'ALPHAFOLD3_READY_{VERSION}', 'w').close()\n", " print(f'Packages installed ({alphafold3.__file__}).')\n", "\n", "# tokamax's Triton kernels need more shared memory than Ada cards have, so\n", @@ -147,7 +160,7 @@ " os.system(f'(python {_script} {_args} > {STAMP}.log 2>&1 && touch {STAMP}) &')\n", "\n", "# /tmp is wiped with the VM, so a fresh session recompiles (~53 s); Drive survives.\n", - "CACHE_DIR = '/tmp/af3_cache'\n", + "CACHE_DIR = os.path.join(__import__('tempfile').gettempdir(), 'af3_cache')\n", "if persist_cache_to_drive:\n", " try:\n", " from google.colab import drive\n", @@ -306,7 +319,7 @@ "jobname = basejob + '_' + hashlib.sha1(flat.encode()).hexdigest()[:5]\n", "\n", "# Input JSON goes to a temp dir; ALL results land in ONE folder: af3_output//\n", - "INPUT_DIR = '/tmp/af3_inputs'\n", + "INPUT_DIR = os.path.join(__import__('tempfile').gettempdir(), 'af3_inputs')\n", "OUTPUT_DIR = 'af3_output'\n", "job_dir = f'{OUTPUT_DIR}/{jobname}'\n", "\n", @@ -323,7 +336,89 @@ " json.dump(fold_input, f, indent=2)\n", "\n", "print(f'Job \"{jobname}\" -> results will be written to {job_dir}/')\n", - "fold_input\n" + "fold_input\n", + "\n", + "from colabfold.rgi.notebook import chain_table\n", + "print('Chain IDs for restraints (residue numbering starts at 1 in each chain):')\n", + "for row in chain_table(fold_input):\n", + " print(row)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "rgi-restraints" + }, + "outputs": [], + "source": [ + "#@title RGI: Restraint-Guided Inference for manipulating molecular structures! (optional)\n", + "\n", + "#@markdown Fill out this form before selecting **Runtime → Run all**. To run standard ColabFold without RGI, leave `use_rgi` unchecked. \n", + "#@markdown Leave any unused restraint fields blank.\n", + "\n", + "#@markdown For more details, see [this preprint](https://www.biorxiv.org/content/10.64898/2026.10.05.756905v1).\n", + "#@markdown For practical examples, including how to specify multiple restraints, refer to the [RGI-toolkit guide](https://github.com/cddlab/rgi_toolkit/blob/main/docs/colabfold.md).\n", + "\n", + "use_rgi = False #@param {type:\"boolean\"}\n", + "#@markdown **conformer restraints** — enter chain IDs such as `B,C`, or `ligands` for all ligand chains. Leave empty to\n", + "#@markdown disable. \n", + "\n", + "#@markdown **To ensure that the ligand stereochemistry matches exactly what is specified in the input (ligand SMILES or CCD), \n", + "#@markdown enable `use_rgi` and select the ligand chains under conformer restraints.**\n", + "\n", + "#@markdown Note: If you apply restraints between protein and ligand molecules, be sure to select the corresponding protein chains as well. \n", + "target_chains = \"\" #@param {type:\"string\"}\n", + "# rename to conformer_chains = target_chains\n", + "conformer_chains = target_chains # rgi_toolkit's config_from_fields reads this name\n", + "#@markdown **distance restraints** — enter two atom selections and their target distance in Å. Example: `chain A and resid\n", + "#@markdown 1 to 10` and `chain A and resid 40 to 50`, with `target_distance = 25`.\n", + "\n", + "\n", + "#@markdown All fields accept the full RGI-toolkit selection language. Residues start at 1 within\n", + "#@markdown each chain. The input cell prints the chain IDs. Distance and angle use atom-group centroids.\n", + "#@markdown Please refer to this manual for the specification method: [Atom selection DSL](https://github.com/cddlab/rgi_toolkit/blob/main/docs/config.md#distance_restraints_config-list)\n", + "distance_atom_selection1 = \"\" #@param {type:\"string\"}\n", + "distance_atom_selection2 = \"\" #@param {type:\"string\"}\n", + "target_distance = 3.5 #@param {type:\"raw\"}\n", + "#@markdown **angle restraints** — enter three atom selections and the target in degrees. Selection 2 is the vertex. Leave\n", + "#@markdown all three empty to disable.\n", + "angle_atom_selection1 = \"\" #@param {type:\"string\"}\n", + "angle_atom_selection2 = \"\" #@param {type:\"string\"}\n", + "angle_atom_selection3 = \"\" #@param {type:\"string\"}\n", + "target_angle = 90 #@param {type:\"raw\"}\n", + "#@markdown **custom** — set named selections and an energy expression, e.g. `harmonic(distance(A, B), 25.0)`.\n", + "#@markdown Leave the energy empty to disable.\n", + "\n", + "#@markdown Example: \"custom_selections:\", `{\"A\": \"chain A and resid 1 to 10\", \"B\": \"chain A and resid 40 to 50\"}`\n", + "custom_selections = {} #@param {type:\"raw\"}\n", + "custom_energy = \"\" #@param {type:\"string\"}\n", + "#@markdown **RMSD** — upload a reference using Colab's **Files → Upload** and enter its filename in either\n", + "#@markdown `ref_pdb` or `ref_cif`. Leave both filenames empty to disable. Empty target/ref selections mean the\n", + "#@markdown whole structure; e.g. `chain A and name CA` selects C-alpha atoms.\n", + "\n", + "#@markdown See also this manual: [rmsd restraints config](https://github.com/cddlab/rgi_toolkit/blob/main/docs/config.md#rmsd_restraints_config-list)\n", + "ref_pdb = \"\" #@param {type:\"string\"}\n", + "ref_cif = \"\" #@param {type:\"string\"}\n", + "atom_selection_target = \"\" #@param {type:\"string\"}\n", + "atom_selection_ref = \"\" #@param {type:\"string\"}\n", + "target_rmsd = 0 #@param {type:\"raw\"}\n", + "\n", + "#@markdown **Advanced (optional)** — an ordinary RGI-toolkit dict for additional restraints, flat-bottomed\n", + "#@markdown penalties, weights, activation windows or conformer settings. Lists append to the fields above. Use\n", + "#@markdown `{\"config_path\": \"restraints.yaml\"}` to read a file. [Examples and native config\n", + "#@markdown reference](https://github.com/cddlab/rgi_toolkit/blob/main/docs/colabfold.md).\n", + "restraints_config = '{}' #@param {type:\"string\"}\n", + "#@markdown After changing this form, rerun this cell and the prediction cell, or use **Run all**. Settings are\n", + "#@markdown saved in the notebook itself.\n", + "\n", + "#@markdown **Model:** choose an RGI-supported model such as `boltz2` in Install dependencies. OpenBind-0 supports RGI. AF2 and IntelliFold2 are vanilla only.\n", + "\n", + "# A raw field left empty is a SyntaxError in Colab, so this one is a string: empty means {}.\n", + "import ast as _ast\n", + "if isinstance(restraints_config, str):\n", + " restraints_config = _ast.literal_eval(restraints_config) if restraints_config.strip() else {}\n" ] }, { @@ -356,6 +451,29 @@ "num_recycles = max(1, int(num_recycles))\n", "num_diffusion_samples = max(1, int(num_diffusion_samples))\n", "\n", + "# Read the standard Colab fields before starting inference.\n", + "if globals().get('use_rgi', False):\n", + " from colabfold.rgi import require_supported_model\n", + " require_supported_model(model)\n", + " subprocess.run(['uv', 'pip', 'install', '--python', sys.executable, _toolkit_source, 'jaxopt'], check=True)\n", + " from rgi_toolkit.notebook_colab import config_from_fields\n", + " rgi_config, rgi_conformer_chains = config_from_fields(globals())\n", + "else:\n", + " use_rgi = False\n", + " rgi_config, rgi_conformer_chains = None, ''\n", + "\n", + "# Rebuild the input even when this cell is rerun after switching RGI off.\n", + "from colabfold.rgi.notebook import name_job, prepare_input\n", + "fold_input = prepare_input(fold_input, model, use_rgi=use_rgi, config=rgi_config,\n", + " conformer_chains=rgi_conformer_chains)\n", + "jobname = name_job(fold_input, model, base=basejob, num_recycles=num_recycles,\n", + " num_diffusion_samples=num_diffusion_samples)\n", + "fold_input['name'] = jobname\n", + "job_dir = os.path.join(OUTPUT_DIR, jobname)\n", + "json_path = os.path.join(INPUT_DIR, f'{jobname}.json')\n", + "with open(json_path, 'w') as _f:\n", + " json.dump(fold_input, _f, indent=2)\n", + "\n", "os.makedirs(OUTPUT_DIR, exist_ok=True)\n", "\n", "# Re-run policy (one folder per job, no timestamped duplicates):\n", @@ -417,7 +535,7 @@ "print(f'Model: {model}')\n", "\n", "cmd = [\n", - " 'python', 'run_alphafold.py',\n", + " sys.executable, *(['-m', 'colabfold.rgi'] if use_rgi else []), 'run_alphafold.py',\n", " f'--json_path={json_path}',\n", " f'--model={model}',\n", " '--norun_data_pipeline',\n", @@ -439,11 +557,12 @@ "if IS_AF3 or IS_AF2:\n", " cmd.append(f'--model_dir={AF2_DIR if IS_AF2 else NATIVE_DIR}')\n", "\n", - "cmd = ' '.join(cmd)\n", + "import shlex\n", + "print('Command:', shlex.join(cmd))\n", "if run_it:\n", " print(cmd)\n", " # Popen rather than `!`: streams the output and gives an exit status.\n", - " _p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,\n", + " _p = subprocess.Popen(cmd, stdout=subprocess.PIPE,\n", " stderr=subprocess.STDOUT, text=True, bufsize=1)\n", " for _line in _p.stdout:\n", " print(_line, end='')\n", @@ -456,7 +575,20 @@ " print(f'\\nDone -> {job_dir}/ ({len(_cifs)} structures, '\n", " f'{time.time() - _T0:.0f} s)')\n", "else:\n", - " print(f'Skipping: results already exist in {job_dir}/ (set on_existing=overwrite to recompute).')\n" + " print(f'Skipping: results already exist in {job_dir}/ (set on_existing=overwrite to recompute).')\n", + "\n", + "if use_rgi:\n", + " _report_path = os.path.join(job_dir, 'rgi_report.json')\n", + " if not os.path.isfile(_report_path):\n", + " raise RuntimeError('RGI report is missing; the run cannot be verified. Re-run with on_existing=overwrite.')\n", + " _rgi_report = json.load(open(_report_path))\n", + " from colabfold.citations import write_rgi_citation\n", + " write_rgi_citation(job_dir)\n", + " for _record in _rgi_report['seeds']:\n", + " print('RGI seed', _record['seed'], 'restraint counts:', _record['inventory'])\n", + " for _distance in _record['distances']:\n", + " print('Measured distances (Angstrom):', _distance['distances_angstrom'],\n", + " 'target/bounds:', _distance['lower'], _distance['upper'])\n" ] }, { @@ -560,7 +692,7 @@ "token_chain_ids = conf.get('token_chain_ids', []) # PAE is per-TOKEN\n", "pae = np.array(conf.get('pae', []), dtype=float)\n", "\n", - "# \u2500\u2500 Summary (ipTM is None for single-chain jobs \u2014 guard before formatting) \u2500\n", + "# ── Summary (ipTM is None for single-chain jobs — guard before formatting) ─\n", "def fmt(v):\n", " return f'{v:.3f}' if isinstance(v, (int, float)) else 'n/a'\n", "\n", @@ -572,17 +704,17 @@ "print('=' * 38)\n", "print(f'Mean pLDDT : {fmt(mean_plddt)}')\n", "print(f'pTM : {fmt(summ.get(\"ptm\"))}')\n", - "print(f'ipTM : {fmt(iptm)}' + (' (single chain \u2014 no interface)' if iptm is None else ''))\n", + "print(f'ipTM : {fmt(iptm)}' + (' (single chain — no interface)' if iptm is None else ''))\n", "print(f'Ranking score : {fmt(summ.get(\"ranking_score\"))}')\n", "print('=' * 38)\n", "\n", - "# \u2500\u2500 Plots \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n", + "# ── Plots ───────────────────────────────────────────────────\n", "has_pae = pae.ndim == 2 and pae.size > 0\n", "ncols = 2 if has_pae else 1\n", "fig, axes = plt.subplots(1, ncols, figsize=(13 if has_pae else 6.5, 4))\n", "axes = np.atleast_1d(axes)\n", "\n", - "# pLDDT per residue \u2014 a line (coloured per chain when there is more than one)\n", + "# pLDDT per residue — a line (coloured per chain when there is more than one)\n", "ax = axes[0]\n", "x = np.arange(len(plddts))\n", "xmax = max(len(plddts) - 1, 1)\n", @@ -613,7 +745,7 @@ "if has_pae:\n", " ax = axes[1]\n", " im = ax.imshow(pae, cmap='bwr', vmin=0, vmax=30, interpolation='nearest')\n", - " plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label='PAE (\u00c5)')\n", + " plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label='PAE (Å)')\n", " if token_chain_ids:\n", " for b in [i for i in range(1, len(token_chain_ids)) if token_chain_ids[i] != token_chain_ids[i-1]]:\n", " ax.axhline(b - 0.5, c='black', lw=0.8)\n", @@ -638,11 +770,17 @@ "source": [ "#@title Download results\n", "from google.colab import files\n", - "import os\n", - "\n", - "results_zip = f'{jobname}.result.zip'\n", - "os.system(f'zip -r {results_zip} {OUTPUT_DIR}/{jobname}')\n", - "files.download(results_zip)\n" + "import subprocess\n", + "\n", + "OUTPUT_DIR = globals().get(\"OUTPUT_DIR\", \"af3_output\")\n", + "results_zip = f\"{jobname}.result.zip\"\n", + "subprocess.run(\n", + " [\"zip\", \"-r\", f\"../{results_zip}\", jobname],\n", + " cwd=OUTPUT_DIR,\n", + " check=True\n", + ")\n", + "files.download(results_zip)\n", + "\n" ] }, { @@ -655,7 +793,7 @@ "# Instructions \n", "\n", "**Quick start:** pick a **model** in the install cell, fill in the sequence(s), then\n", - "**Runtime \u2192 Run all**. The first run downloads that model's weights; later runs reuse\n", + "**Runtime → Run all**. The first run downloads that model's weights; later runs reuse\n", "them, and each model has its own cache so switching back is instant.\n", "\n", "---\n", @@ -665,22 +803,24 @@ "Ported weights come from [sokrypton/af3-any-model](https://huggingface.co/sokrypton/af3-any-model).\n", "\"Tower\" is a protein language model fetched separately on first use.\n", "\n", - "| model | weights | licence | notes |\n", - "|---|---|---|---|\n", - "| `openbind0` | [OpenFold3 v0.5.0 \"OpenBind\"](https://github.com/aqlaboratory/openfold-3/releases/tag/v0.5.0) (AlQuraishi Lab) | Apache-2.0 | The current release, and the default here. |\n", - "| `openfold3` | [OpenFold3 preview-2](https://github.com/aqlaboratory/openfold) (AlQuraishi Lab) | Apache-2.0 | The earlier preview, kept because earlier results used it. |\n", - "| `boltz2` | [Boltz-2](https://github.com/jwohlwend/boltz) (Wohlwend et al.) | MIT | Strong on ligands; keeps a modified residue as one token. |\n", - "| `protenix2` | [Protenix-v2](https://github.com/bytedance/Protenix) (ByteDance) | Apache-2.0 | The widest trunk here (pair 256), so the slowest. |\n", - "| `rosettafold3` | [RoseTTAFold3](https://github.com/RosettaCommons/foundry) (RosettaCommons) | BSD-3-Clause | Carries chirality features; handles D-amino acids. |\n", - "| `chai1` | [chai-1](https://github.com/chaidiscovery/chai-lab) (Chai Discovery) | Apache-2.0 | Folds from ESM2 3B, fetched and run automatically. |\n", - "| `intellifold2` | [IntelliFold-v2](https://huggingface.co/intelligenAI/intellifold) (IntelliGen-AI) | Apache-2.0 | Widened channels (pair 512), largest ported download. |\n", - "| `opendde` | [OpenDDE](https://huggingface.co/aurekaresearch/OpenDDE) (Aureka Research) | Apache-2.0 | Runs its diffusion on an expanded structural-token set. |\n", - "| `esmfold2` | [ESMFold2](https://huggingface.co/biohub/ESMFold2) (Chan Zuckerberg Biohub) | MIT | Folds from ESM-C instead of an MSA \u2014 single sequence, no search. |\n", - "| `esmfold2_lm600m` | ESMFold2, 600M tower | MIT | No confidence head. |\n", - "| `esmfold2_lm300m` | ESMFold2, 300M tower | MIT | No confidence head. |\n", - "| `af2_ptm` | AlphaFold 2 monomer pTM (DeepMind) | CC BY 4.0 | **Protein only** \u2014 a ligand or nucleotide raises rather than quietly folding the rest. Templates use the model_1/model_2 parameter sets. |\n", - "| `af2_multimer` | AlphaFold 2 multimer v3 (DeepMind) | CC BY 4.0 | Protein only, as above. |\n", - "| `alphafold3` | Google DeepMind's own parameters | [AF3 terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md) | DeepMind's public release, downloaded on first use (~1 GB). Its terms govern the weights and the outputs; run_alphafold prints them at startup. |\n", + "**RGI**: **Yes** supports `use_rgi = True`; **No** requires `use_rgi = False`.\n", + "\n", + "| model | weights | licence | RGI | notes |\n", + "|---|---|---|---|---|\n", + "| `openbind0` | [OpenFold3 v0.5.0 \"OpenBind\"](https://github.com/aqlaboratory/openfold-3/releases/tag/v0.5.0) (AlQuraishi Lab) | Apache-2.0 | Yes | The current OpenFold3 release, with optional RGI. |\n", + "| `openfold3` | [OpenFold3 preview-2](https://github.com/aqlaboratory/openfold) (AlQuraishi Lab) | Apache-2.0 | Yes | The earlier preview, kept because earlier results used it. |\n", + "| `boltz2` | [Boltz-2](https://github.com/jwohlwend/boltz) (Wohlwend et al.) | MIT | Yes | Strong on ligands; keeps a modified residue as one token. |\n", + "| `protenix2` | [Protenix-v2](https://github.com/bytedance/Protenix) (ByteDance) | Apache-2.0 | Yes | The widest trunk here (pair 256), so the slowest. |\n", + "| `rosettafold3` | [RoseTTAFold3](https://github.com/RosettaCommons/foundry) (RosettaCommons) | BSD-3-Clause | Yes | Carries chirality features; handles D-amino acids. |\n", + "| `chai1` | [chai-1](https://github.com/chaidiscovery/chai-lab) (Chai Discovery) | Apache-2.0 | Yes | Folds from ESM2 3B, fetched and run automatically. |\n", + "| `intellifold2` | [IntelliFold-v2](https://huggingface.co/intelligenAI/intellifold) (IntelliGen-AI) | Apache-2.0 | No | Widened channels (pair 512), largest ported download. |\n", + "| `opendde` | [OpenDDE](https://huggingface.co/aurekaresearch/OpenDDE) (Aureka Research) | Apache-2.0 | Yes | Runs its diffusion on an expanded structural-token set. |\n", + "| `esmfold2` | [ESMFold2](https://huggingface.co/biohub/ESMFold2) (Chan Zuckerberg Biohub) | MIT | Yes | Folds from ESM-C instead of an MSA — single sequence, no search. |\n", + "| `esmfold2_lm600m` | ESMFold2, 600M tower | MIT | Yes | No confidence head. |\n", + "| `esmfold2_lm300m` | ESMFold2, 300M tower | MIT | Yes | No confidence head. |\n", + "| `af2_ptm` | AlphaFold 2 monomer pTM (DeepMind) | CC BY 4.0 | No | **Protein only** — a ligand or nucleotide raises rather than quietly folding the rest. Templates use the model_1/model_2 parameter sets. |\n", + "| `af2_multimer` | AlphaFold 2 multimer v3 (DeepMind) | CC BY 4.0 | No | Protein only, as above. |\n", + "| `alphafold3` | Google DeepMind's own parameters | [AF3 terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md) | Yes | DeepMind's public release, downloaded on first use (~1 GB). Its terms govern the weights and the outputs; run_alphafold prints them at startup. |\n", "---\n", "\n", "## Input\n", @@ -695,37 +835,57 @@ "| **ligand_ccd** | ligand(s) by PDB CCD code | `ATP:MG:HEM` |\n", "| **ligand_smiles** | ligand(s) by SMILES | `CC(=O)Oc1ccccc1C(=O)O` |\n", "\n", - "Mix boxes freely to build a complex. Chain IDs A, B, C\u2026 follow AlphaFold 3's canonical\n", - "order (protein \u2192 RNA \u2192 DNA \u2192 ligand). Identical protein sequences are merged, so\n", + "Mix boxes freely to build a complex. Chain IDs A, B, C… follow AlphaFold 3's canonical\n", + "order (protein → RNA → DNA → ligand). Identical protein sequences are merged, so\n", "`SEQ:SEQ` is a homodimer. Sequences and CCD codes are upper-cased; **SMILES are left\n", - "as typed**. Whitespace and extra colons are forgiven (`SEQ1::::SEQ2` = `SEQ1:SEQ2`) \u2014\n", + "as typed**. Whitespace and extra colons are forgiven (`SEQ1::::SEQ2` = `SEQ1:SEQ2`) —\n", "which is also why an atom-mapped SMILES containing `:` needs a raw AF3 JSON instead.\n", "\n", "**seeds**: comma-separated, one prediction each (`1,2,3`). Junk and duplicates are\n", "dropped. **msa_mode**: `mmseqs2_server` queries the public\n", - "[ColabFold](https://colabfold.mmseqs.com/) API (protein only \u2014 RNA/DNA always run\n", + "[ColabFold](https://colabfold.mmseqs.com/) API (protein only — RNA/DNA always run\n", "MSA-free); `single_sequence` skips it, faster and less accurate.\n", "\n", "## Output\n", "\n", "| file | contents |\n", "|---|---|\n", - "| `*.cif` | Best-ranked structure. B-factor = pLDDT (0\u2013100). |\n", + "| `*.cif` | Best-ranked structure. B-factor = pLDDT (0–100). |\n", "| `*_confidences.json` | Per-residue pLDDT, PAE matrix, contact probabilities. |\n", "| `*_summary_confidences.json` | Mean pLDDT, pTM, ipTM, ranking score. |\n", - "| `*_ranking_scores.csv` | Every seed \u00d7 sample combination. |\n", + "| `*_ranking_scores.csv` | Every seed × sample combination. |\n", "| `seed-N_sample-M/` | One directory per prediction. |\n", "| `TERMS_OF_USE.md` | The licence for whichever weights you ran. |\n", "\n", - "pLDDT above 90 is very high, 70\u201390 reliable backbone, 50\u201370 doubtful, below 50 likely\n", + "pLDDT above 90 is very high, 70–90 reliable backbone, 50–70 doubtful, below 50 likely\n", "disordered or wrong. Lower PAE means two residues are confidently placed *relative to\n", "each other*, which is what to read for an interface. ipTM above 0.8 is a well-defined\n", - "complex interface, and is `n/a` for a single chain \u2014 there is no interface to score.\n", + "complex interface, and is `n/a` for a single chain — there is no interface to score.\n", + "\n", + "### RGI citation\n", + "\n", + "If you use RGI, please cite:\n", + "\n", + "Hori, T., Moriwaki, Y. & Ishitani, R. (2026). *RGI-Toolkit: Differentiable Restraints\n", + "for Controllable Biomolecular Structure Prediction.* bioRxiv, preprint, version 1.\n", + "[doi:10.64898/2026.10.05.756905](https://doi.org/10.64898/2026.10.05.756905).\n", + "\n", + "```bibtex\n", + "@article{hori2026rgi,\n", + " author = {Hori, Tatsuki and Moriwaki, Yoshitaka and Ishitani, Ryuichiro},\n", + " title = {{RGI-Toolkit}: Differentiable Restraints for Controllable Biomolecular Structure Prediction},\n", + " journal = {bioRxiv},\n", + " year = {2026},\n", + " doi = {10.64898/2026.10.05.756905},\n", + " url = {https://www.biorxiv.org/content/10.64898/2026.10.05.756905v1},\n", + " note = {Preprint, version 1}\n", + "}\n", + "```\n", "\n", "## Troubleshooting\n", "\n", - "**OOM**: shorter sequence, or a larger GPU (`Runtime \u2192 Change runtime type`).\n", - "**MSA server timeout**: the public server is rate-limited \u2014 retry, or use\n", + "**OOM**: shorter sequence, or a larger GPU (`Runtime → Change runtime type`).\n", + "**MSA server timeout**: the public server is rate-limited — retry, or use\n", "`single_sequence`. **Download popup blocked**: disable your ad blocker.\n", "\n", "## Licence\n", @@ -747,8 +907,8 @@ "accelerator": "GPU", "colab": { "gpuType": "T4", - "provenance": [], - "include_colab_link": true + "include_colab_link": true, + "provenance": [] }, "kernelspec": { "display_name": "Python 3 (ipykernel)", @@ -770,4 +930,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/colabfold/citations.py b/colabfold/citations.py index 7610b175d..53d70a5c1 100644 --- a/colabfold/citations.py +++ b/colabfold/citations.py @@ -1,9 +1,19 @@ import logging +import re from pathlib import Path logger = logging.getLogger(__name__) citations = { + "hori2026rgi": """@article{hori2026rgi, + author = {Hori, Tatsuki and Moriwaki, Yoshitaka and Ishitani, Ryuichiro}, + title = {{RGI-Toolkit}: Differentiable Restraints for Controllable Biomolecular Structure Prediction}, + journal = {bioRxiv}, + year = {2026}, + doi = {10.64898/2026.10.05.756905}, + url = {https://www.biorxiv.org/content/10.64898/2026.10.05.756905v1}, + note = {Preprint, version 1} +}""", "Mirdita2021": """@article{Mirdita2022, author= {Mirdita, Milot and Schütze, Konstantin and Moriwaki, Yoshitaka and Heo, Lim and Ovchinnikov, Sergey and Steinegger, Martin }, doi = {10.1038/s41592-022-01488-1}, @@ -157,4 +167,16 @@ def write_bibtex( writer.write("\n") logger.info(f"Found {len(to_cite)} citations for tools or databases") - return bibtex_file \ No newline at end of file + return bibtex_file + + +def write_rgi_citation(result_dir: str | Path) -> Path: + """Append the RGI paper once, preserving existing output citations.""" + bibtex_file = Path(result_dir) / "cite.bibtex" + existing = bibtex_file.read_text(encoding="utf-8") if bibtex_file.exists() else "" + if not re.search(r"@\w+\s*\{\s*hori2026rgi\s*,", existing, re.IGNORECASE): + with bibtex_file.open("a", encoding="utf-8") as writer: + if existing and not existing.endswith("\n"): + writer.write("\n") + writer.write(citations["hori2026rgi"] + "\n") + return bibtex_file diff --git a/colabfold/rgi/__init__.py b/colabfold/rgi/__init__.py new file mode 100644 index 000000000..8b76a22cd --- /dev/null +++ b/colabfold/rgi/__init__.py @@ -0,0 +1,26 @@ +"""Optional RGI integration for the ColabFold diffusion notebooks.""" + +SUPPORTED_MODELS = frozenset( + { + "alphafold3", + "openbind0", + "openfold3", + "boltz2", + "protenix2", + "rosettafold3", + "chai1", + "opendde", + "esmfold2", + "esmfold2_lm600m", + "esmfold2_lm300m", + } +) +ALPHAFOLD3_COLABFOLD_VERSION = "3.1.11" + + +def require_supported_model(model): + if model not in SUPPORTED_MODELS: + raise ValueError( + f"RGI is not supported for {model!r}. Choose " + f"{', '.join(sorted(SUPPORTED_MODELS))}, or turn use_rgi off." + ) diff --git a/colabfold/rgi/__main__.py b/colabfold/rgi/__main__.py new file mode 100644 index 000000000..ccc599313 --- /dev/null +++ b/colabfold/rgi/__main__.py @@ -0,0 +1,27 @@ +"""Run the pinned upstream CLI with optional input-driven RGI hooks.""" + +import importlib.util +import sys +from pathlib import Path + + +def main(): + if len(sys.argv) < 2 or sys.argv[1] in {"-h", "--help"}: + print("Usage: python -m colabfold.rgi run_alphafold.py [upstream flags]") + return + script = Path(sys.argv[1]).resolve() + if not script.is_file(): + raise FileNotFoundError(script) + sys.argv = [str(script), *sys.argv[2:]] + spec = importlib.util.spec_from_file_location("_rgi_colabfold_runner", script) + runner = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = runner + spec.loader.exec_module(runner) + from colabfold.rgi.runtime import install + + install(runner) + runner.app.run(runner.main) + + +if __name__ == "__main__": + main() diff --git a/colabfold/rgi/boltz.py b/colabfold/rgi/boltz.py new file mode 100644 index 000000000..f5e0ba223 --- /dev/null +++ b/colabfold/rgi/boltz.py @@ -0,0 +1,52 @@ +"""Preserve legacy notebook molecules and MSAs in Boltz's RGI-capable YAML.""" + +import copy + + +def notebook_input( + entries, *, config=None, conformer_chains="ligands", single_sequence=False +): + sequences = [] + ligands = set() + for header, sequence in entries: + fields = header.lstrip(">").split("|") + chain, kind = fields[:2] + molecule = {"id": chain} + if kind == "protein": + molecule.update( + sequence=sequence, msa="empty" if single_sequence else fields[2] + ) + elif kind.lower() in {"dna", "rna"}: + kind = kind.lower() + molecule["sequence"] = sequence + elif kind in {"smiles", "ccd"}: + molecule[kind] = sequence + kind = "ligand" + ligands.add(chain) + else: + raise ValueError(f"Unsupported molecule type {kind!r}.") + sequences.append({kind: molecule}) + raw = {"version": 1, "sequences": sequences} + if config is not None: + from rgi_toolkit.config import RestraintsConfig + + RestraintsConfig.from_dict(config) + raw["restraints_config"] = copy.deepcopy(config) + if config.get("conformer_restraints_config") is not None: + chosen = ( + ligands + if conformer_chains == "ligands" + else { + chain.strip() + for chain in conformer_chains.split(",") + if chain.strip() + } + ) + known = {next(iter(entity.values()))["id"] for entity in sequences} + if not chosen or not chosen <= known: + raise ValueError("Choose existing conformer chains or add a ligand.") + for entity in sequences: + molecule = next(iter(entity.values())) + if molecule["id"] in chosen: + molecule["conformer_restraints"] = True + return raw diff --git a/colabfold/rgi/notebook.py b/colabfold/rgi/notebook.py new file mode 100644 index 000000000..2f337b1d3 --- /dev/null +++ b/colabfold/rgi/notebook.py @@ -0,0 +1,95 @@ +"""Input preparation shared by the notebook forms and headless E2E runs.""" + +from __future__ import annotations + +import copy +import hashlib +import json +import re + +from colabfold.rgi import require_supported_model + + +def chain_table(fold_input): + rows = [] + for entity in fold_input["sequences"]: + kind, molecule = next(iter(entity.items())) + ids = molecule["id"] + for chain in [ids] if isinstance(ids, str) else ids: + rows.append( + { + "chain": chain, + "type": kind, + "residues": len(molecule["sequence"]) + if "sequence" in molecule + else None, + } + ) + return rows + + +def prepare_input( + fold_input, model, *, use_rgi=False, config=None, conformer_chains="ligands" +): + """Return a fresh input; disabling RGI removes every RGI opt-in.""" + raw = copy.deepcopy(fold_input) + raw.pop("restraints_config", None) + entities = [] + for entity in raw["sequences"]: + kind, molecule = next(iter(entity.items())) + molecule.pop("conformer_restraints", None) + ids = molecule["id"] + for chain in [ids] if isinstance(ids, str) else ids: + entities.append({kind: dict(molecule, id=chain)}) + raw["sequences"] = entities + if model.startswith("esmfold2"): + proteins = [row for row in chain_table(raw) if row["type"] == "protein"] + if len(proteins) != 1: + raise ValueError( + "ESMFold2 language-model inputs need exactly one protein chain. " + "Use one sequence or choose another predictor for a protein complex." + ) + if not use_rgi: + return raw + require_supported_model(model) + from rgi_toolkit.config import RestraintsConfig + + if not config: + raise ValueError("RGI is enabled but no restraints were configured.") + RestraintsConfig.from_dict(config) + if config.get("conformer_restraints_config") is not None: + table = chain_table(raw) + if conformer_chains.strip() == "ligands": + chosen = {row["chain"] for row in table if row["type"] == "ligand"} + else: + chosen = { + item.strip() for item in conformer_chains.split(",") if item.strip() + } + known = {row["chain"] for row in table} + if not chosen or not chosen <= known: + raise ValueError( + "Choose existing conformer chains from the table, or add a ligand." + ) + for entity in entities: + molecule = next(iter(entity.values())) + if molecule["id"] in chosen: + molecule["conformer_restraints"] = True + raw["restraints_config"] = copy.deepcopy(config) + return raw + + +def name_job(fold_input, model, *, base, num_recycles, num_diffusion_samples): + """Keep changed model, sampling, and restraint settings out of stale results.""" + material = {key: value for key, value in fold_input.items() if key != "name"} + material.update( + model=model, + num_recycles=num_recycles, + num_diffusion_samples=num_diffusion_samples, + ) + digest = hashlib.sha256(json.dumps(material, sort_keys=True).encode()).hexdigest()[ + :10 + ] + base = re.sub(r"\W+", "", base).lower() or "job" + if "restraints_config" in fold_input: + return f"{base}_{model}_rgi_{digest}" + return f"{base}_{model}_{digest}" diff --git a/colabfold/rgi/runtime.py b/colabfold/rgi/runtime.py new file mode 100644 index 000000000..0d1f273a7 --- /dev/null +++ b/colabfold/rgi/runtime.py @@ -0,0 +1,274 @@ +"""Scoped glue between the ColabFold JAX runner and the shared RGI engine.""" + +from __future__ import annotations + +import copy +import functools +import inspect +import json +from contextlib import contextmanager +from pathlib import Path + +from colabfold.citations import write_rgi_citation +from colabfold.rgi import ALPHAFOLD3_COLABFOLD_VERSION, require_supported_model + + +@contextmanager +def _replace(obj, name, value): + original = getattr(obj, name) + setattr(obj, name, value) + try: + yield + finally: + setattr(obj, name, original) + + +def split_input(raw, base_dir=None): + """Strip only RGI extensions before calling the unchanged upstream parser.""" + from rgi_toolkit.config import RestraintsConfig, resolve_restraints_config + + clean = copy.deepcopy(raw) + config = clean.pop("restraints_config", None) + opted = {} + for entity in clean.get("sequences", []): + for molecule in entity.values(): + flag = molecule.pop("conformer_restraints", False) + if not isinstance(flag, bool): + raise TypeError("conformer_restraints must be true or false.") + ids = molecule.get("id", []) + for chain in [ids] if isinstance(ids, str) else ids: + opted[chain] = flag + if config is not None: + config = resolve_restraints_config(config, base_dir=base_dir) + RestraintsConfig.from_dict(config) + elif any(opted.values()): + raise ValueError("Conformer opt-in requires restraints_config.") + return clean, config, opted + + +def build_adapter(fold_input, example, opted): + """Resolve predictor chemistry; reuse the existing framework-free AF3 adapter.""" + from alphafold3.common import folding_input + from alphafold3.constants import decoded_ccd, residue_names + from alphafold3.data.tools import rdkit_utils + from rdkit import Chem + from rgi_toolkit.alphafold3.adapter import AF3RestraintAdapter + + mols = [] + ccd = None + for chain in fold_input.chains: + if not opted.get(chain.id) or not isinstance(chain, folding_input.Ligand): + continue + if chain.smiles is not None: + mol = Chem.MolFromSmiles(chain.smiles) + if mol is None: + raise ValueError(f"Invalid ligand SMILES for chain {chain.id}.") + mols.append((chain.id, mol, True)) + else: + if len(chain.ccd_ids or []) != 1: + raise ValueError( + "Ligand geometry currently needs one CCD per ligand chain." + ) + if ccd is None: + ccd = decoded_ccd.get_ccd(user_ccd=fold_input.user_ccd) + mol = rdkit_utils.mol_from_ccd_cif( + ccd[chain.ccd_ids[0]], sort_alphabetically=False, remove_hydrogens=True + ) + mols.append((chain.id, mol, False)) + mapping = {chain.id: i + 1 for i, chain in enumerate(fold_input.chains)} + return AF3RestraintAdapter( + example, + mapping, + residue_names.POLYMER_TYPES, + mols, + {mapping[chain]: flag for chain, flag in opted.items() if chain in mapping}, + ) + + +def _guided_apply(runner, model_runner): + import haiku as hk + import jax + from alphafold3.model.network import diffusion_head + + from colabfold.rgi.sampler import coordinate_minimizer, make_guided_sampler + + original = diffusion_head.sample + + @hk.transform + def forward(batch, minimize, gather): + sampler = make_guided_sampler(original, coordinate_minimizer(minimize, gather)) + with _replace(diffusion_head, "sample", sampler): + return runner.model.Model(model_runner._model_config)( + batch, use_dropout=model_runner._use_dropout + ) + + model_runner._preinit_tokamax_context() + return forward.apply if runner._NOJIT.value else jax.jit(forward.apply) + + +def install(runner): + """Install input and per-job hooks in one CLI process, preserving vanilla.""" + import importlib.metadata + + if ( + importlib.metadata.version("alphafold3-colabfold") + != ALPHAFOLD3_COLABFOLD_VERSION + ): + raise RuntimeError( + "RGI requires alphafold3-colabfold==3.1.11; re-run installation." + ) + + input_class = runner.folding_input.Input + original_parse = input_class.from_json + original_predict = runner.predict_structure + original_write_input = runner.write_fold_input_json + original_write_outputs = runner.write_outputs + jobs, reports = {}, {} + + @classmethod + def parse(cls, json_str, json_path=None): + raw = json.loads(json_str) + clean, config, opted = split_input( + raw, Path(json_path).parent if json_path is not None else None + ) + result = original_parse(json.dumps(clean), json_path) + settings = (config, opted) + if result.name in jobs and jobs[result.name] != settings: + raise ValueError("Give jobs with different RGI settings different names.") + jobs[result.name] = settings + return result + + def write_input(fold_input, output_dir): + original_write_input(fold_input, output_dir) + # A failed rerun must never leave the previous run's success report. + (Path(output_dir) / "rgi_report.json").unlink(missing_ok=True) + config, opted = jobs.get(fold_input.name, (None, {})) + if config is not None: + path = Path(output_dir) / f"{fold_input.sanitised_name()}_data.json" + raw = json.loads(path.read_text()) + raw["restraints_config"] = config + entities = [] + for entity in raw["sequences"]: + kind = next(iter(entity)) + molecule = next(iter(entity.values())) + ids = molecule["id"] + ids = [ids] if isinstance(ids, str) else ids + for chain in ids: + item = dict(molecule, id=chain) + if opted.get(chain): + item["conformer_restraints"] = True + entities.append({kind: item}) + raw["sequences"] = entities + path.write_text(json.dumps(raw, indent=2) + "\n") + + def predict(fold_input, model_runner, **kwargs): + import jax.numpy as jnp + import numpy as np + from rgi_toolkit.combined import CombinedRestraints + from rgi_toolkit.notebook import distance_report, restraint_inventory + + config, opted = jobs.get(fold_input.name, (None, {})) + if config is None: + return original_predict(fold_input, model_runner, **kwargs) + require_supported_model(model_runner.model_name) + records = [] + original_inference = model_runner.run_inference + if not hasattr(model_runner, "_rgi_apply"): + model_runner._rgi_apply = _guided_apply(runner, model_runner) + + def inference(example, rng_key): + restraints = CombinedRestraints() + adapter = build_adapter(fold_input, example, opted) + restraints.setup(adapter, config=config) + if not restraints.is_active(): + raise ValueError( + "RGI built no active restraints. " + "Check selections and ligand opt-ins." + ) + if config.get("conformer_restraints_config") is not None and not any( + opted.values() + ): + raise ValueError( + "Ligand/polymer geometry requires " + "conformer_restraints: true on a chain." + ) + inventory = restraint_inventory(restraints) + print( + "RGI built restraints:", + json.dumps(inventory, sort_keys=True), + flush=True, + ) + minimize = restraints.get_minimizer() + gather = example.get("structbook/residue_atom_gather") + if gather is not None: + gather = np.asarray(gather) + mask = np.asarray(example["ref_mask"], bool) + if gather.shape != mask.shape or np.any(mask & (gather < 0)): + raise ValueError( + "OpenDDE structural atom map does not cover the RGI atoms." + ) + valid = gather[gather >= 0] + if len(np.unique(valid)) != len(valid): + raise ValueError( + "OpenDDE structural atom map contains duplicate atoms." + ) + gather = jnp.asarray(gather) + apply = model_runner._rgi_apply + had_model = "_model" in vars(model_runner) + cached_model = vars(model_runner).get("_model") + model_runner._model = lambda key, batch: apply( + model_runner.model_params, key, batch, minimize, gather + ) + try: + result = original_inference(example, rng_key) + finally: + if had_model: + model_runner._model = cached_model + else: + del model_runner._model + coords = np.asarray(result["diffusion_samples"]["atom_positions"]) + flat = coords.reshape(coords.shape[0], -1, 3) + if not np.isfinite(flat).all(): + raise RuntimeError( + "The predictor produced non-finite atom coordinates." + ) + restraints.finalize(jnp.asarray(flat), 0) + records.append( + { + "inventory": inventory, + "finite": True, + "distances": distance_report(restraints, flat), + } + ) + print( + "RGI measured distances:", + json.dumps(records[-1]["distances"]), + flush=True, + ) + return result + + with _replace(model_runner, "run_inference", inference): + results = original_predict(fold_input, model_runner, **kwargs) + for record, result in zip(records, results, strict=True): + record["seed"] = result.seed + reports[fold_input.sanitised_name()] = { + "model": model_runner.model_name, + "restraints_config": config, + "seeds": records, + } + return results + + @functools.wraps(original_write_outputs) + def write_outputs(*args, **kwargs): + original_write_outputs(*args, **kwargs) + bound = inspect.signature(original_write_outputs).bind(*args, **kwargs) + report = reports.pop(bound.arguments["job_name"], None) + if report is not None: + path = Path(bound.arguments["output_dir"]) / "rgi_report.json" + path.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n") + write_rgi_citation(path.parent) + + input_class.from_json = parse + runner.predict_structure = predict + runner.write_fold_input_json = write_input + runner.write_outputs = write_outputs diff --git a/colabfold/rgi/sampler.py b/colabfold/rgi/sampler.py new file mode 100644 index 000000000..5672e3d78 --- /dev/null +++ b/colabfold/rgi/sampler.py @@ -0,0 +1,87 @@ +"""Add checked hook points to the pinned upstream sampler in memory. + +The original sampler owns all schedules, random keys, Euler/Heun steps and +alignment. No installed file is modified, and vanilla keeps the original function. +""" + +from __future__ import annotations + +import hashlib +import inspect +from functools import lru_cache + + +@lru_cache(maxsize=4) +def _instrument(original): + source = inspect.getsource(original) + digest = hashlib.sha256(source.encode()).hexdigest() + if digest != EXPECTED_SAMPLE_SHA256: + raise RuntimeError( + "The installed diffusion sampler differs from the tested release. " + "Re-run the install cell with alphafold3-colabfold==3.1.11. " + f"Sampler SHA256: {digest}" + ) + changes = ( + ( + "noise_level, noise_level_prev = step_noise", + "noise_level, noise_level_prev, rgi_step = step_noise", + ), + ( + "positions_denoised = denoising_step(positions_noisy, t_hat)", + ( + "positions_denoised = denoising_step(positions_noisy, t_hat)\n" + " positions_denoised = _rgi_minimize(" + "positions_denoised, noise_level_prev, rgi_step)" + ), + ), + ( + "denoising_step(positions_out, noise_level)", + ( + "_rgi_minimize(denoising_step(positions_out, noise_level), " + "noise_level_prev, rgi_step)" + ), + ), + ( + "(noise_levels[1:], noise_levels[:-1]),", + ( + "(noise_levels[1:], noise_levels[:-1], " + "jnp.arange(noise_levels.shape[0] - 1))," + ), + ), + ) + for old, new in changes: + if source.count(old) != 1: + raise RuntimeError(f"Cannot locate exactly one RGI sampler hook: {old}") + source = source.replace(old, new) + return compile(source, inspect.getsourcefile(original), "exec") + + +def make_guided_sampler(original, minimize): + namespace = dict(original.__globals__, _rgi_minimize=minimize) + exec(_instrument(original), namespace) # noqa: S102 - SHA256-checked upstream source. + return namespace[original.__name__] + + +def coordinate_minimizer(minimize, gather=None): + """Bridge dense AF3 or OpenDDE structural-token coordinates to the engine.""" + import jax.numpy as jnp + + def apply(coords, sigma, step): + flat = coords.reshape(-1, 3) + if gather is None: + return minimize(flat, sigma, step).reshape(coords.shape) + # OpenDDE diffuses structural subtokens; RGI retains per-residue numbering. + # Invalid padded entries scatter out of bounds with mode='drop'. + indices = gather.reshape(-1) + valid = indices >= 0 + residue = jnp.where(valid[:, None], flat[jnp.maximum(indices, 0)], 0) + corrected = minimize(residue, sigma, step) + destinations = jnp.where(valid, indices, flat.shape[0]) + return flat.at[destinations].set(corrected, mode="drop").reshape(coords.shape) + + return apply + + +EXPECTED_SAMPLE_SHA256 = ( + "d238e26b0b7f4f86efaa121c1ea2b2279b673c97dc9b7ed332f4ceffbe43ef1c" +) diff --git a/docs/rgi.md b/docs/rgi.md new file mode 100644 index 000000000..8cc27f74e --- /dev/null +++ b/docs/rgi.md @@ -0,0 +1,34 @@ +# RGI in the standard Colab form + +Open [ColabFold2 preview](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/ColabFold2_preview.ipynb#scrollTo=rgi-restraints). +**RGI: Making restraint-guided protein folding inference accessible to all! (optional)** +is a standard Colab form cell after **Input sequences**. +Its fields are visible and editable before execution. + +1. Choose an RGI-supported model such as **openbind0** in **Install dependencies**. +2. Fill in **Input sequences**. +3. In the **RGI** form, turn on **use_rgi** and enter the selections and targets. +4. Use **Runtime → Run all**. + +For a 25 Å distance, set `distance_atom_selection1` to `chain A and resid 1 to 10`, +`distance_atom_selection2` to `chain A and resid 40 to 50`, and `target_distance` to +`25`. Use chains and residues from your own input. Leave unused restraint fields empty. + +The same form provides conformer, angle, custom and RMSD. Lists in the fields create +multiple restraints; `restraints_config` accepts additional native toolkit entries +or `{"config_path": "restraints.yaml"}`. See the +[full guide and examples](https://github.com/cddlab/rgi_toolkit/blob/main/docs/colabfold.md). + +OpenBind-0 uses the same RGI sampler hook as OpenFold3 in the shared JAX runner, +while retaining its own model configuration and checkpoint. It is available in +both ColabFold2 preview and AlphaFold3 / OpenFold3. + +Leave **use_rgi** off for vanilla. AF2 and IntelliFold2 are vanilla only. +Vanilla runs use the upstream predictor without the RGI wrapper or RGI dependency. +Any restraint configuration and chain opt-ins left from an earlier run are removed. +The notebooks retain their default models and sampling settings. +After edits, rerun the RGI cell and prediction, or use **Run all**. In Boltz-1 the form +precedes installation; rerun installation too when changing **use_rgi**. +RGI results include the RGI-Toolkit paper in `cite.bibtex`, which is also included +in the downloaded results ZIP. Existing citations are preserved without duplicating +the RGI entry on reruns. diff --git a/scripts/rgi_batch_e2e.py b/scripts/rgi_batch_e2e.py new file mode 100644 index 000000000..caffaa4f1 --- /dev/null +++ b/scripts/rgi_batch_e2e.py @@ -0,0 +1,112 @@ +"""Verify vanilla/guided transitions and config isolation in one real JAX process.""" + +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path + +from rgi_e2e import measure, prepare_runtime +from rgi_toolkit.notebook import make_config + + +def run(work): + work = Path(work).resolve() + work.mkdir(parents=True, exist_ok=True) + runner = prepare_runtime(work) + inputs = work / "inputs" + inputs.mkdir(exist_ok=True) + targets = { + "a_vanilla": None, + "b_distance25": 25, + "c_distance35": 35, + "d_custom30": 30, + "e_vanilla": None, + } + for name, target in targets.items(): + data = { + "name": name, + "dialect": "alphafold3", + "version": 1, + "modelSeeds": [42, 43], + "sequences": [ + { + "protein": { + "id": chain, + "sequence": "ACDEFGHIK", + "templates": [], + "unpairedMsa": ">query\nACDEFGHIK\n", + "pairedMsa": "", + } + } + for chain in ("A", "B") + ], + } + if name == "d_custom30": + data["restraints_config"] = { + "verbose": True, + "custom_restraints_config": [ + { + "name": "centroid_target", + "selections": {"A": "chain A", "B": "chain B"}, + "energy": "(distance(A, B) - 30)**2", + } + ], + } + elif target is not None: + data["restraints_config"] = make_config( + { + "distance_restraints_config": [ + { + "atom_selection1": "chain A", + "atom_selection2": "chain B", + "harmonic": {"target_distance": target}, + } + ] + } + ) + (inputs / f"{name}.json").write_text(json.dumps(data, indent=2)) + os.environ.setdefault("XLA_FLAGS", "--xla_gpu_enable_triton_gemm=false") + os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") + command = [ + sys.executable, + "-m", + "colabfold.rgi", + str(runner), + f"--input_dir={inputs}", + "--model=boltz2", + "--norun_data_pipeline", + f"--output_dir={work / 'outputs'}", + "--force_output_dir", + "--flash_attention_implementation=xla", + "--num_recycles=1", + "--num_diffusion_samples=2", + "--buckets=32", + ] + with (work / "batch.log").open("w") as log: + subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True) + summary = {} + for name, target in targets.items(): + job = work / "outputs" / name + cifs = sorted(job.glob("seed-*_sample-*/*.cif")) + assert len(cifs) == 4, (name, cifs) + values = [measure(path) for path in cifs] + if target is not None: + assert all(abs(value - target) < 0.5 for value in values), (name, values) + report = json.loads((job / "rgi_report.json").read_text()) + assert [record["seed"] for record in report["seeds"]] == [42, 43] + term = "custom" if name == "d_custom30" else "distance" + assert all(record["inventory"][term] == 1 for record in report["seeds"]) + else: + assert not (job / "rgi_report.json").exists() + summary[name] = values + assert summary["a_vanilla"] == summary["e_vanilla"], summary + (work / "summary.json").write_text(json.dumps(summary, indent=2)) + print("BATCH E2E PASSED", json.dumps(summary), flush=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--work-dir", default=".cache/rgi-e2e/batch") + run(parser.parse_args().work_dir) diff --git a/scripts/rgi_boltz1_e2e.py b/scripts/rgi_boltz1_e2e.py new file mode 100644 index 000000000..346a1bec4 --- /dev/null +++ b/scripts/rgi_boltz1_e2e.py @@ -0,0 +1,91 @@ +"""Exercise native Boltz-1 using the exact notebook input builder.""" + +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path + +import yaml +from rgi_e2e import measure +from rgi_toolkit.notebook import make_config + +from colabfold.rgi.boltz import notebook_input + + +def run(args): + work = Path(args.work_dir).resolve() + work.mkdir(parents=True, exist_ok=True) + summary = [] + entries = [(f">{chain}|protein|empty", "ACDEFGHIK") for chain in ("A", "B")] + for arm in ("vanilla", "rgi"): + config = ( + None + if arm == "vanilla" + else make_config( + { + "verbose": True, + "distance_restraints_config": [ + { + "atom_selection1": "chain A", + "atom_selection2": "chain B", + "harmonic": {"target_distance": 25}, + } + ], + } + ) + ) + data = notebook_input(entries, config=config, single_sequence=True) + path = work / f"boltz1_{arm}.yaml" + path.write_text(yaml.safe_dump(data, sort_keys=False)) + cmd = [ + args.python, + "-m", + "boltz.main", + "predict", + str(path), + "--model", + "boltz1", + "--out_dir", + str(work), + "--recycling_steps", + "1", + "--diffusion_samples", + "1", + "--seed", + "42", + "--override", + "--num_workers", + "0", + ] + log_path = work / f"boltz1_{arm}.log" + with log_path.open("w") as log: + proc = subprocess.run( + cmd, stdout=log, stderr=subprocess.STDOUT, check=False + ) + if proc.returncode: + raise RuntimeError(log_path.read_text()[-12000:]) + cifs = list( + (work / f"boltz_results_{path.stem}" / "predictions" / path.stem).glob( + "*.cif" + ) + ) + assert len(cifs) == 1, cifs + distance = measure(cifs[0]) + if arm == "rgi": + assert "distances=1" in log_path.read_text() + assert abs(distance - 25) < 0.5, distance + summary.append({"model": "boltz1", "arm": arm, "distance_angstrom": distance}) + print(json.dumps(summary[-1]), flush=True) + assert abs(summary[0]["distance_angstrom"] - summary[1]["distance_angstrom"]) > 1 + (work / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--python", default=os.environ.get("RGI_BOLTZ_PYTHON", sys.executable) + ) + parser.add_argument("--work-dir", default=".cache/rgi-e2e/boltz1") + run(parser.parse_args()) diff --git a/scripts/rgi_e2e.py b/scripts/rgi_e2e.py new file mode 100644 index 000000000..990cc9ae0 --- /dev/null +++ b/scripts/rgi_e2e.py @@ -0,0 +1,221 @@ +"""Run paired real predictions and validate written structures and RGI evidence. + +Run with an environment containing the notebook dependencies: + uv run --no-project python scripts/rgi_e2e.py --models boltz2 +Use --work-dir to keep downloaded parameters and outputs between invocations. +""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import os +import subprocess +import sys +import urllib.request +from pathlib import Path + + +def prepare_runtime(work): + from alphafold3.constants import ccd_fetch + + runner = work / "run_alphafold.py" + if not runner.exists(): + urllib.request.urlretrieve( + "https://raw.githubusercontent.com/sokrypton/alphafold3/v3.1.11/run_alphafold.py", + runner, + ) + root = Path(importlib.metadata.distribution("alphafold3-colabfold").locate_file("")) + conv = root / "alphafold3/constants/converters" + conv.mkdir(parents=True, exist_ok=True) + if not (conv / "ccd.pickle").exists(): + ccd_fetch.write_pickles( + ccd_fetch.codes_for_input(extra=["ATP"]), + str(conv / "ccd.pickle"), + str(conv / "chemical_component_sets.pickle"), + libcifpp_dir=str(root / "share/libcifpp"), + ) + return runner + + +def measure(path): + import gemmi + import numpy as np + + structure = gemmi.read_structure(str(path)) + groups = {} + for chain in structure[0]: + xyz = np.asarray( + [ + [atom.pos.x, atom.pos.y, atom.pos.z] + for residue in chain + for atom in residue + ] + ) + if not len(xyz) or not np.isfinite(xyz).all(): + raise AssertionError(f"Empty or non-finite chain in {path.name}") + groups[chain.name] = xyz.mean(axis=0) + if "B" not in groups: + residues = list(structure[0]["A"]) + groups = { + key: np.asarray( + [[a.pos.x, a.pos.y, a.pos.z] for r in subset for a in r] + ).mean(axis=0) + for key, subset in (("A", residues[:9]), ("B", residues[9:])) + } + return float(np.linalg.norm(groups["A"] - groups["B"])) + + +def run(args): + import numpy as np + from rgi_toolkit.notebook import make_config + + work = Path(args.work_dir).resolve() + work.mkdir(parents=True, exist_ok=True) + runner = prepare_runtime(work) + os.environ.setdefault("XLA_FLAGS", "--xla_gpu_enable_triton_gemm=false") + os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") + os.environ.setdefault("AF3_WEIGHTS_DIR", str(work / "weights")) + summary = [] + for model in args.models: + distances = {} + for arm in args.arms: + name = f"{model}_{arm}" + sequence = "ACDEFGHIK" + raw = { + "name": name, + "dialect": "alphafold3", + "version": 1, + "modelSeeds": [42], + "sequences": [ + { + "protein": { + "id": chain, + "sequence": sequence, + "templates": [], + "unpairedMsa": f">query\n{sequence}\n", + "pairedMsa": "", + } + } + for chain in ("A", "B") + ], + } + if model.startswith("esmfold2"): + protein = raw["sequences"][0]["protein"] + protein["sequence"] = sequence * 2 + protein["unpairedMsa"] = f">query\n{sequence * 2}\n" + raw["sequences"] = raw["sequences"][:1] + if args.fixture == "combined": + raw["sequences"].append( + {"ligand": {"id": "C", "smiles": "O=C(O)/C=C/C(=O)O"}} + ) + if arm != "vanilla": + target = 25.0 if arm == "rgi" else 35.0 + config = { + "verbose": True, + "distance_restraints_config": [ + { + "atom_selection1": "chain A and resid 1 to 9" + if model.startswith("esmfold2") + else "chain A", + "atom_selection2": "chain A and resid 10 to 18" + if model.startswith("esmfold2") + else "chain B", + "harmonic": {"target_distance": target}, + } + ], + } + if args.fixture == "combined": + config["conformer_restraints_config"] = {"plane": {"weight": 1.0}} + raw["sequences"][-1]["ligand"]["conformer_restraints"] = True + raw["restraints_config"] = make_config(config) + path = work / f"{name}.json" + path.write_text(json.dumps(raw, indent=2) + "\n") + out = work / "outputs" + command = [ + sys.executable, + *([] if arm == "vanilla" else ["-m", "colabfold.rgi"]), + str(runner), + f"--json_path={path}", + f"--model={model}", + "--norun_data_pipeline", + f"--output_dir={out}", + "--force_output_dir", + "--flash_attention_implementation=xla", + "--num_recycles=1", + "--num_diffusion_samples=1", + "--buckets=32", + "--weights_precision=fp32" + if model == "alphafold3" + else "--weights_precision=int8", + f"--cache_dir={work / 'compile_cache'}", + ] + if model == "chai1" or model.startswith("esmfold2"): + command.append("--use_esm_embeddings") + if model == "alphafold3": + weights = work / "af3_weights" + weights.mkdir(exist_ok=True) + if not (weights / "af3.bin.zst").exists(): + urllib.request.urlretrieve( + "https://storage.googleapis.com/alphafold3/af3.bin.zst", + weights / "af3.bin.zst", + ) + command.append(f"--model_dir={weights}") + log_path = work / f"{name}.log" + print(f"Running {model} / {arm}; log: {log_path.name}", flush=True) + with log_path.open("w") as log: + proc = subprocess.run( + command, stdout=log, stderr=subprocess.STDOUT, check=False + ) + if proc.returncode: + print(log_path.read_text()[-12000:], flush=True) + raise RuntimeError(f"{model}/{arm} exited {proc.returncode}") + job = out / name + cifs = sorted(job.glob("seed-*_sample-*/*.cif")) + if not cifs: + raise AssertionError(f"No structures for {model}/{arm}") + values = [measure(cif) for cif in cifs] + distances[arm] = values + report = job / "rgi_report.json" + if arm != "vanilla": + evidence = json.loads(report.read_text()) + assert evidence["seeds"][0]["inventory"]["distance"] == 1 + if args.fixture == "combined": + inventory = evidence["seeds"][0]["inventory"] + assert all( + inventory[key] > 0 + for key in ("bond", "angle", "plane", "cistrans") + ), inventory + assert all(abs(value - target) < 0.5 for value in values), values + elif report.exists(): + raise AssertionError("Vanilla must not produce an RGI report") + summary.append( + { + "model": model, + "arm": arm, + "samples": len(cifs), + "distance_angstrom": values, + } + ) + (work / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary[-1]), flush=True) + if "rgi" in distances and "vanilla" in distances: + assert not np.allclose(distances["rgi"], distances["vanilla"], atol=1.0) + print("E2E PASSED", json.dumps(summary), flush=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--models", nargs="+", default=["boltz2"]) + parser.add_argument( + "--arms", + nargs="+", + choices=["vanilla", "rgi", "rgi_other"], + default=["vanilla", "rgi"], + ) + parser.add_argument("--work-dir", default=".cache/rgi-e2e/run") + parser.add_argument( + "--fixture", choices=["distance", "combined"], default="distance" + ) + run(parser.parse_args()) diff --git a/tests/test_citations.py b/tests/test_citations.py new file mode 100644 index 000000000..445731845 --- /dev/null +++ b/tests/test_citations.py @@ -0,0 +1,132 @@ +import importlib.metadata +import json +import shutil +import sys +from pathlib import Path +from types import SimpleNamespace +from zipfile import ZipFile + +import pytest + +from colabfold.citations import write_bibtex, write_rgi_citation + + +@pytest.mark.parametrize("existing", [False, True]) +def test_rgi_citation_preserves_existing_bibliography_on_rerun(tmp_path, existing): + original = b"" + if existing: + path = write_bibtex("alphafold2_ptm", True, True, True, True, tmp_path) + original = path.read_bytes() + assert b"hori2026rgi" not in original + + path = write_rgi_citation(tmp_path) + first = path.read_bytes() + assert first.startswith(original) + assert first.count(b"@article{hori2026rgi,") == 1 + assert b"doi = {10.64898/2026.10.05.756905}" in first + assert write_rgi_citation(tmp_path) == path + assert path.read_bytes() == first + + +def test_rgi_citation_recognizes_existing_bibtex_key(tmp_path): + path = tmp_path / "cite.bibtex" + original = "@ARTICLE { hori2026rgi ,\n title = {Existing citation}\n}\n" + path.write_text(original, encoding="utf-8") + write_rgi_citation(tmp_path) + assert path.read_text(encoding="utf-8") == original + + +def test_runtime_output_citation_is_scoped_to_guided_jobs(tmp_path, monkeypatch): + pytest.importorskip("jax") + pytest.importorskip("rgi_toolkit") + from colabfold.rgi import ALPHAFOLD3_COLABFOLD_VERSION, runtime + + class Input: + @classmethod + def from_json(cls, json_str, json_path=None): + result = cls() + result.name = json.loads(json_str)["name"] + return result + + def sanitised_name(self): + return self.name + + def write_outputs(results, output_dir, job_name): + output_dir.mkdir() + write_bibtex("", False, False, False, False, output_dir) + + runner = SimpleNamespace( + folding_input=SimpleNamespace(Input=Input), + predict_structure=lambda *args, **kwargs: [], + write_fold_input_json=lambda *args: None, + write_outputs=write_outputs, + ) + model_runner = SimpleNamespace( + model_name="boltz2", _rgi_apply=None, run_inference=lambda *args: None + ) + original_version = importlib.metadata.version + monkeypatch.setattr( + importlib.metadata, + "version", + lambda name: ( + ALPHAFOLD3_COLABFOLD_VERSION + if name == "alphafold3-colabfold" + else original_version(name) + ), + ) + runtime.install(runner) + for name, guided in (("guided", True), ("vanilla", False)): + raw = {"name": name} + if guided: + raw["restraints_config"] = {} + fold_input = Input.from_json(json.dumps(raw)) + results = runner.predict_structure(fold_input, model_runner) + output_dir = tmp_path / name + runner.write_outputs(results, output_dir=output_dir, job_name=name) + bibliography = (output_dir / "cite.bibtex").read_text() + assert "@article{Mirdita2022," in bibliography + assert ("@article{hori2026rgi," in bibliography) is guided + assert (output_dir / "rgi_report.json").exists() is guided + + +@pytest.mark.parametrize( + ("notebook", "download_id", "output_dir"), + [ + ("ColabFold2_preview.ipynb", "download", "outputs/test"), + ("AlphaFold3_of3.ipynb", "download", "outputs/test"), + ("Boltz1.ipynb", "jdSBSTOpaULF", "test/boltz_results_test"), + ], +) +def test_downloaded_results_include_rgi_citation( + notebook, download_id, output_dir, tmp_path, monkeypatch +): + if notebook != "Boltz1.ipynb" and shutil.which("zip") is None: + pytest.skip("The AF3 notebook download cell requires zip.") + root = Path(__file__).resolve().parents[1] + cells = json.loads((root / notebook).read_text())["cells"] + source = next( + "".join(cell["source"]) + for cell in cells + if cell["metadata"]["id"] == download_id + ) + monkeypatch.chdir(tmp_path) + output = Path(output_dir) + output.mkdir(parents=True) + citation = write_rgi_citation(output).read_bytes() + input_yaml = tmp_path / "test.yaml" + input_yaml.write_text("version: 1\n") + downloaded = [] + monkeypatch.setitem( + sys.modules, + "google.colab", + SimpleNamespace(files=SimpleNamespace(download=downloaded.append)), + ) + exec( + source, + {"jobname": "test", "OUTPUT_DIR": "outputs", "input_yaml": input_yaml}, + ) + assert len(downloaded) == 1 + with ZipFile(downloaded[0]) as archive: + names = [name for name in archive.namelist() if name.endswith("cite.bibtex")] + assert len(names) == 1 + assert archive.read(names[0]) == citation diff --git a/tests/test_notebook_vanilla.py b/tests/test_notebook_vanilla.py new file mode 100644 index 000000000..6554b1bd2 --- /dev/null +++ b/tests/test_notebook_vanilla.py @@ -0,0 +1,271 @@ +"""Keep the notebook's vanilla command independent of RGI state and dependencies.""" + +import builtins +import copy +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +@pytest.mark.parametrize( + ("notebook", "model", "extra_args"), + [ + ("ColabFold2_preview.ipynb", "openbind0", []), + ("ColabFold2_preview.ipynb", "af2_ptm", ["--model_dir=af2_weights"]), + ("ColabFold2_preview.ipynb", "chai1", ["--use_esm_embeddings"]), + ("ColabFold2_preview.ipynb", "esmfold2", ["--use_esm_embeddings"]), + ("ColabFold2_preview.ipynb", "esmfold2_lm300m", ["--use_esm_embeddings"]), + ("AlphaFold3_of3.ipynb", "alphafold3", ["--model_dir=af3_weights"]), + ], +) +def test_vanilla_prediction_ignores_previous_rgi_state( + notebook, model, extra_args, tmp_path, monkeypatch +): + root = Path(__file__).resolve().parents[1] + cells = json.loads((root / notebook).read_text())["cells"] + source = next( + "".join(cell["source"]) for cell in cells if cell["metadata"]["id"] == "run" + ) + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("AF3_NB_OVERRIDES", "{}") + Path("inputs").mkdir() + original = { + "name": "previous_run", + "modelSeeds": [42], + "sequences": [ + { + "protein": { + "id": "A", + "sequence": "ACDEFGHIK", + "unpairedMsa": ">query\nACDEFGHIK\n", + "pairedMsa": "", + "templates": [], + } + } + ], + } + guided = copy.deepcopy(original) + guided["restraints_config"] = {"invalid_previous_config": True} + guided["sequences"][0]["protein"]["conformer_restraints"] = True + namespace = { + "model": model, + "sys": sys, + "json": json, + "fold_input": guided, + "basejob": "test", + "INPUT_DIR": "inputs", + "OUTPUT_DIR": "outputs", + "CACHE_DIR": "cache", + "AF2_DIR": "af2_weights", + "NATIVE_DIR": "af3_weights", + "IS_AF2": model.startswith("af2_"), + "IS_AF3": model == "alphafold3", + "msa_mode": "single_sequence", + "on_existing": "overwrite", + "use_rgi": False, + "restraints_config": "invalid", + "ref_pdb": "missing.pdb", + } + commands = [] + + def run(command, **kwargs): + assert command[0] == "nvidia-smi", command + return SimpleNamespace(stdout="7.5\n", returncode=0) + + def popen(command, **kwargs): + commands.append(command) + output = Path(namespace["job_dir"]) + output.mkdir(parents=True) + (output / "result.cif").touch() + return SimpleNamespace(stdout=iter([]), wait=lambda: 0) + + real_import = builtins.__import__ + + def without_rgi_toolkit(name, *args, **kwargs): + if name == "rgi_toolkit" or name.startswith("rgi_toolkit."): + raise AssertionError("Vanilla prediction must not import rgi-toolkit") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setattr(subprocess, "Popen", popen) + monkeypatch.setattr(builtins, "__import__", without_rgi_toolkit) + exec(source, namespace) + + assert len(commands) == 1 + command = commands[0] + assert command[:2] == [sys.executable, "run_alphafold.py"] + assert f"--model={model}" in command + assert "--num_recycles=10" in command + assert "--num_diffusion_samples=5" in command + assert [ + arg + for arg in command + if arg == "--use_esm_embeddings" or arg.startswith("--model_dir=") + ] == extra_args + assert not any("rgi" in argument for argument in command[1:]) + written = json.loads(Path(namespace["json_path"]).read_text()) + assert written["sequences"] == original["sequences"] + assert written["modelSeeds"] == original["modelSeeds"] + assert "restraints_config" not in written + assert namespace["rgi_config"] is None + assert not (Path(namespace["job_dir"]) / "cite.bibtex").exists() + + +@pytest.mark.parametrize( + ("rgi_cached", "produce_structure"), [(False, True), (True, True), (True, False)] +) +def test_native_boltz_vanilla_run_all_switches_back_to_upstream( + rgi_cached, produce_structure, tmp_path, monkeypatch +): + yaml = pytest.importorskip( + "yaml", reason="The native Boltz notebook requires PyYAML." + ) + + root = Path(__file__).resolve().parents[1] + cells = json.loads((root / "Boltz1.ipynb").read_text())["cells"] + sources = {cell["metadata"]["id"]: "".join(cell["source"]) for cell in cells} + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(shutil, "which", lambda name: name) + Path("test").mkdir() + Path("test/query.a3m").write_text(">A\nACDEFGHIK\n") + stale_output = Path("test/boltz_results_test") + (stale_output / "processed").mkdir(parents=True) + (stale_output / "processed/manifest.json").write_text("{}") + stale_predictions = stale_output / "predictions/test" + stale_predictions.mkdir(parents=True) + (stale_predictions / "stale.cif").touch() + if rgi_cached: + (stale_output / ".rgi-enabled").touch() + (stale_output / "cite.bibtex").write_text("@article{hori2026rgi}\n") + namespace = { + "os": os, + "use_rgi": False, + "boltz_python": ".cache/boltz-rgi/bin/python", + "restraints_config": "invalid", + "ref_pdb": "missing.pdb", + "fasta_entries": [(">A|protein", "ACDEFGHIK"), (">D|smiles", "CCO")], + "msa_mode": "single_sequence", + "jobname": "test", + } + commands = [] + + def run(command, **kwargs): + commands.append(command) + if "predict" in command: + if rgi_cached: + assert not stale_output.exists() + else: + assert (stale_output / "processed/manifest.json").is_file() + assert (stale_predictions / "stale.cif").is_file() + assert Path("test/query.a3m").read_text() == ">A\nACDEFGHIK\n" + if produce_structure: + output = stale_output / "predictions/test" + output.mkdir(parents=True, exist_ok=True) + (output / "result.cif").touch() + return SimpleNamespace(returncode=0) + + real_import = builtins.__import__ + + def without_rgi_toolkit(name, *args, **kwargs): + if name == "rgi_toolkit" or name.startswith("rgi_toolkit."): + raise AssertionError("Vanilla prediction must not import rgi-toolkit") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setattr(builtins, "__import__", without_rgi_toolkit) + exec(sources["4eXNO1JJHYrB"], namespace) + if produce_structure: + exec(sources["bgaBXxXtIAu9"], namespace) + else: + with pytest.raises( + RuntimeError, match="Prediction produced no structure files" + ): + exec(sources["bgaBXxXtIAu9"], namespace) + assert not (stale_output / ".rgi-enabled").exists() + assert not (stale_output / "cite.bibtex").exists() + + assert any(command[-1] == "boltz==2.2.1" for command in commands) + assert not any( + "boltz_restr" in arg or "rgi_toolkit" in arg + for command in commands + for arg in command + ) + command = commands[-1] + assert command[:4] == [ + ".cache/boltz-vanilla/bin/python", + "-m", + "boltz.main", + "predict", + ] + assert command[command.index("--model") + 1] == "boltz1" + assert command[command.index("--recycling_steps") + 1] == "3" + assert command[command.index("--diffusion_samples") + 1] == "1" + written = yaml.safe_load(Path("test/test.yaml").read_text()) + assert "restraints_config" not in written + assert all( + "conformer_restraints" not in next(iter(entity.values())) + for entity in written["sequences"] + ) + + +def test_native_boltz_rgi_reprocesses_when_target_changes(tmp_path, monkeypatch): + pytest.importorskip("yaml", reason="The native Boltz notebook requires PyYAML.") + root = Path(__file__).resolve().parents[1] + cells = json.loads((root / "Boltz1.ipynb").read_text())["cells"] + source = next( + "".join(cell["source"]) + for cell in cells + if cell["metadata"]["id"] == "bgaBXxXtIAu9" + ) + monkeypatch.chdir(tmp_path) + output = Path("test/boltz_results_test") + (output / "processed").mkdir(parents=True) + (output / "processed/manifest.json").write_text("{}") + namespace = { + "Path": Path, + "subprocess": subprocess, + "use_rgi": True, + "boltz_python": "boltz-rgi-python", + "fasta_entries": [(">A|protein", "ACDEFGHIK")], + "msa_mode": "single_sequence", + "jobname": "test", + } + + def config_from_fields(fields): + return {"target_distance": fields["target_distance"]}, "" + + monkeypatch.setitem( + sys.modules, + "rgi_toolkit.notebook_colab", + SimpleNamespace(config_from_fields=config_from_fields), + ) + monkeypatch.setattr( + "colabfold.rgi.boltz.notebook_input", + lambda entries, *, config, **kwargs: {"restraints_config": config}, + ) + targets = [] + + def run(command, **kwargs): + assert (output / ".rgi-enabled").is_file() + assert not (output / "processed").exists() + targets.append(namespace["rgi_config"]["target_distance"]) + (output / "processed").mkdir() + (output / "processed/manifest.json").write_text("{}") + prediction = output / "predictions/test" + prediction.mkdir(parents=True) + (prediction / "result.cif").touch() + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(subprocess, "run", run) + for target in (25, 35): + namespace["target_distance"] = target + exec(source, namespace) + bibliography = (output / "cite.bibtex").read_text() + assert bibliography.count("@article{hori2026rgi,") == 1 + assert targets == [25, 35] diff --git a/tests/test_rgi_notebook.py b/tests/test_rgi_notebook.py new file mode 100644 index 000000000..ca0d40cd7 --- /dev/null +++ b/tests/test_rgi_notebook.py @@ -0,0 +1,271 @@ +"""Input isolation, novice controls, and optional-dependency boundaries.""" + +import copy +import importlib.util +import json +import re +from pathlib import Path + +import pytest + +from colabfold.rgi import SUPPORTED_MODELS, require_supported_model +from colabfold.rgi.notebook import chain_table, name_job, prepare_input +from colabfold.rgi.runtime import split_input + +pytestmark = pytest.mark.skipif( + importlib.util.find_spec("rgi_toolkit") is None, + reason="The optional RGI tests need rgi-toolkit and Python 3.12 or newer.", +) + + +@pytest.fixture +def raw(): + return { + "name": "demo", + "modelSeeds": [1], + "sequences": [ + {"protein": {"id": ["A", "B"], "sequence": "ACDEFG"}}, + {"ligand": {"id": "C", "smiles": "CCO"}}, + ], + } + + +def test_every_supported_model_accepts_optional_rgi(raw): + config = {"conformer_restraints_config": {}} + for model in SUPPORTED_MODELS: + data = copy.deepcopy(raw) + if model.startswith("esmfold2"): + data["sequences"][0]["protein"]["id"] = "A" + guided = prepare_input(data, model, use_rgi=True, config=config) + assert guided["sequences"][-1]["ligand"]["conformer_restraints"] is True + assert "conformer_restraints" not in guided["sequences"][0]["protein"] + vanilla = prepare_input(guided, model, use_rgi=False, config=config) + assert "restraints_config" not in vanilla + assert all( + "conformer_restraints" not in next(iter(x.values())) + for x in vanilla["sequences"] + ) + assert "restraints_config" not in raw + + +def test_unsupported_models_keep_vanilla(raw): + for model in ("af2_ptm", "af2_multimer", "intellifold2", "unknown"): + prepare_input(raw, model) + with pytest.raises(ValueError, match="turn use_rgi off"): + require_supported_model(model) + + +def test_conformer_controls_do_not_silently_do_nothing(raw): + config = {"conformer_restraints_config": {}} + for chains in ("", "Z", "C,Z"): + with pytest.raises(ValueError): + prepare_input( + raw, "boltz2", use_rgi=True, config=config, conformer_chains=chains + ) + raw["sequences"].pop() + with pytest.raises(ValueError): + prepare_input(raw, "boltz2", use_rgi=True, config=config) + + +def test_chain_table_and_partial_homomer_opt_in(raw): + assert [r["chain"] for r in chain_table(raw)] == ["A", "B", "C"] + guided = prepare_input( + raw, + "boltz2", + use_rgi=True, + config={"conformer_restraints_config": {}}, + conformer_chains="B", + ) + assert "conformer_restraints" not in guided["sequences"][0]["protein"] + assert guided["sequences"][1]["protein"]["conformer_restraints"] is True + + +def test_result_names_change_with_effective_settings(raw): + def name(data=raw, model="boltz2", recycles=1): + return name_job( + data, model, base="demo", num_recycles=recycles, num_diffusion_samples=1 + ) + + changed = copy.deepcopy(raw) + changed["restraints_config"] = {"conformer_restraints_config": {}} + assert len({name(), name(changed), name(model="openfold3"), name(recycles=2)}) == 4 + + +def test_result_names_mark_only_rgi_runs(raw): + def name(data): + return name_job( + data, "boltz2", base="demo", num_recycles=1, num_diffusion_samples=1 + ) + + guided = copy.deepcopy(raw) + guided["restraints_config"] = {"conformer_restraints_config": {}} + assert re.fullmatch(r"demo_boltz2_[0-9a-f]{10}", name(raw)) + assert re.fullmatch(r"demo_boltz2_rgi_[0-9a-f]{10}", name(guided)) + + +def test_parser_removes_only_rgi_extensions_and_keeps_flags_separate(raw): + raw = prepare_input( + raw, "boltz2", use_rgi=True, config={"conformer_restraints_config": {}} + ) + clean, config, flags = split_input(raw) + assert config == {"conformer_restraints_config": {}} + assert flags == {"A": False, "B": False, "C": True} + assert "restraints_config" not in clean + assert "conformer_restraints" not in clean["sequences"][-1]["ligand"] + assert "restraints_config" in raw + + +def test_esm_complex_error_is_actionable(raw): + with pytest.raises(ValueError, match="exactly one protein chain"): + prepare_input(raw, "esmfold2") + + +def test_boltz_yaml_preserves_molecules_and_msa_paths(): + from colabfold.rgi.boltz import notebook_input + + entries = [ + (">A|protein|msa/query.a3m", "ACDE"), + (">B|smiles", "CCO"), + (">C|DNA", "ACGT"), + ] + config = {"conformer_restraints_config": {}} + guided = notebook_input(entries, config=config) + assert guided["sequences"][0]["protein"]["msa"] == "msa/query.a3m" + assert guided["sequences"][1]["ligand"]["smiles"] == "CCO" + assert guided["sequences"][1]["ligand"]["conformer_restraints"] is True + assert guided["sequences"][2]["dna"]["sequence"] == "ACGT" + vanilla = notebook_input(entries, single_sequence=True) + assert "restraints_config" not in vanilla + assert vanilla["sequences"][0]["protein"]["msa"] == "empty" + + +def test_vanilla_does_not_import_the_optional_engine(raw, monkeypatch): + import builtins + + original = builtins.__import__ + + def guarded(name, *args, **kwargs): + if name.startswith("rgi_toolkit"): + raise ImportError("Optional engine is deliberately unavailable") + return original(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", guarded) + assert "restraints_config" not in prepare_input(raw, "boltz2") + + +def test_notebooks_have_no_outputs_and_python_cells_compile(): + root = Path(__file__).resolve().parents[1] + for name in ("ColabFold2_preview.ipynb", "AlphaFold3_of3.ipynb", "Boltz1.ipynb"): + notebook = json.loads((root / name).read_text()) + for cell in notebook["cells"]: + if cell["cell_type"] == "code": + assert cell["outputs"] == [] + compile("".join(cell["source"]), name, "exec") + + +@pytest.mark.parametrize( + ("name", "model"), + [ + ("ColabFold2_preview.ipynb", "openbind0"), + ("AlphaFold3_of3.ipynb", "openbind0"), + ("Boltz1.ipynb", "boltz2"), + ], +) +def test_native_form_runs_without_widgets_and_reads_every_edit( + name, model, raw, monkeypatch +): + import subprocess + import sys + + from rgi_toolkit.notebook_colab import COLAB_FORM + + monkeypatch.setattr(subprocess, "run", lambda *args, **kwargs: None) + monkeypatch.setenv("AF3_NB_OVERRIDES", "{}") + root = Path(__file__).resolve().parents[1] + cells = json.loads((root / name).read_text())["cells"] + cell = next(c for c in cells if c.get("metadata", {}).get("id") == "rgi-restraints") + source = "".join(cell["source"]) + assert source.partition("\n")[2].startswith(COLAB_FORM.partition("\n")[2]) + assert "notebook_widgets" not in json.dumps(cells) + assert sum("use_rgi = False #@param" in "".join(c["source"]) for c in cells) == 1 + namespace = {"model": model, "sys": sys, "_toolkit_source": "installed-in-test"} + exec( + source, namespace + ) # The form has no install, input or running-widget prerequisite. + prediction = next( + "".join(c["source"]) + for c in cells + if "# Read the standard Colab fields" in "".join(c["source"]) + ) + prefix = prediction.split( + "# Rebuild the input" if name != "Boltz1.ipynb" else "# Use YAML" + )[0] + exec(prefix, namespace) + assert namespace["rgi_config"] is None + namespace["use_rgi"] = True + with pytest.raises(ValueError, match="RGI is on but no restraints"): + exec(prefix, namespace) + namespace.update( + distance_atom_selection1="chain A and resid 1 to 2", + distance_atom_selection2='["chain A and resid 3", "chain A and resid 4"]', + target_distance=[25, 30], + conformer_chains="C", + ) + exec(prefix, namespace) + config = namespace["rgi_config"] + assert len(config["distance_restraints_config"]) == 2 + guided = prepare_input( + raw, + model, + use_rgi=True, + config=config, + conformer_chains=namespace["rgi_conformer_chains"], + ) + assert guided["sequences"][-1]["ligand"]["conformer_restraints"] is True + namespace["target_distance"] = [15, 20] + exec(prefix, namespace) + assert ( + namespace["rgi_config"]["distance_restraints_config"][1]["harmonic"][ + "target_distance" + ] + == 20 + ) + namespace.update(use_rgi=False, restraints_config="invalid", ref_pdb="missing.pdb") + exec(prefix, namespace) + assert namespace["rgi_config"] is None + assert "restraints_config" not in prepare_input(guided, model, use_rgi=False) + if name != "Boltz1.ipynb": + namespace.update(use_rgi=True, model="af2_ptm") + with pytest.raises(ValueError, match="turn use_rgi off"): + exec(prefix, namespace) + + +def test_preview_preserves_upstream_layout_and_default_model(): + root = Path(__file__).resolve().parents[1] + cells = json.loads((root / "ColabFold2_preview.ipynb").read_text())["cells"] + original = [c for c in cells if c["metadata"]["id"] != "rgi-restraints"] + assert [c["metadata"]["id"] for c in original] == [ + "view-in-github", + "header", + "install", + "input", + "run", + "display3d", + "plots", + "download", + "instructions", + ] + assert [ + "".join(c["source"]).splitlines()[0] + for c in original + if c["cell_type"] == "code" + ] == [ + "#@title Install dependencies (~35 s)", + "#@title Input sequences", + "#@title Run the model", + "#@title Display structures + PAE (py2Dmol)", + "#@title Quality metrics and plots", + "#@title Download results", + ] + assert 'model = "openbind0" #@param' in "".join(original[2]["source"]) + assert "Where to configure RGI" not in "".join(original[1]["source"]) diff --git a/tests/test_rgi_sampler.py b/tests/test_rgi_sampler.py new file mode 100644 index 000000000..f54492303 --- /dev/null +++ b/tests/test_rgi_sampler.py @@ -0,0 +1,87 @@ +"""Exercise the actual pinned sampler under JIT, vmap, and scan.""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +from colabfold.rgi.sampler import coordinate_minimizer, make_guided_sampler + + +def test_structural_token_coordinates_preserve_residue_mapping(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + coords = jnp.arange(18, dtype=jnp.float32).reshape(3, 2, 3) + gather = jnp.array([[4, 1], [3, -1]]) + apply = coordinate_minimizer(lambda x, sigma, step: x + 10, gather) + actual = jax.jit(apply)(coords, 1.0, 0).reshape(-1, 3) + expected = np.asarray(coords).reshape(-1, 3).copy() + expected[[4, 1, 3]] += 10 + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("model", ["openbind0", "chai1", "esmfold2"]) +def test_identity_hook_preserves_upstream_sampling_paths(model): + dh = pytest.importorskip("alphafold3.model.network.diffusion_head") + import haiku as hk + import jax + import jax.numpy as jnp + + config = dh.SampleConfig( + steps=6, num_samples=2, max_sigma=256 if model.startswith("esmfold2") else 0 + ) + batch = SimpleNamespace( + predicted_structure_info=SimpleNamespace(atom_mask=jnp.ones((2, 2))) + ) + gc = SimpleNamespace(model=model) + + def make(sample): + @hk.transform + def forward(): + return sample(lambda x, sigma: x / 2, batch, hk.next_rng_key(), config, gc)[ + "atom_positions" + ] + + return jax.jit(forward.apply) + + original = dh.sample + guided = make_guided_sampler(original, lambda x, sigma, step: x) + key = jax.random.PRNGKey(42) + np.testing.assert_array_equal(make(original)({}, key), make(guided)({}, key)) + assert dh.sample is original + + +def test_schedule_passes_pre_churn_sigma_and_zero_based_steps(): + dh = pytest.importorskip("alphafold3.model.network.diffusion_head") + import haiku as hk + import jax + import jax.numpy as jnp + + seen = [] + + def minimize(x, sigma, step): + jax.debug.callback(lambda s, i: seen.append((float(s), int(i))), sigma, step) + return x + + guided = make_guided_sampler(dh.sample, minimize) + batch = SimpleNamespace( + predicted_structure_info=SimpleNamespace(atom_mask=jnp.ones((2, 2))) + ) + config = dh.SampleConfig(steps=6, num_samples=1) + + @hk.transform + def forward(): + return guided( + lambda x, sigma: x / 2, + batch, + hk.next_rng_key(), + config, + SimpleNamespace(model="openbind0"), + ) + + jax.block_until_ready(jax.jit(forward.apply)({}, jax.random.PRNGKey(42))) + expected = dh.noise_schedule(np.linspace(0, 1, 7, dtype=np.float32))[:-1] + assert sorted(i for s, i in seen) == list(range(6)) + np.testing.assert_allclose( + [s for s, i in sorted(seen, key=lambda x: x[1])], expected + )