From 615a2654b13a2bb31d6c7095a064294837559879 Mon Sep 17 00:00:00 2001 From: Sergey Ovchinnikov Date: Wed, 16 Sep 2026 18:15:59 -0400 Subject: [PATCH 01/11] Created using Colab --- ColabFold2_preview.ipynb | 822 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 822 insertions(+) create mode 100644 ColabFold2_preview.ipynb diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb new file mode 100644 index 000000000..7e1ab5d61 --- /dev/null +++ b/ColabFold2_preview.ipynb @@ -0,0 +1,822 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "id": "header", + "metadata": { + "id": "header" + }, + "source": [ + "# ColabFold2 preview\n", + "\n", + "Predict protein, RNA, DNA and small-molecule structures with [AlphaFold 3](https://www.nature.com/articles/s41586-024-07487-w), running **any** of fourteen models through one implementation. Pick a model in the install cell — the weights download themselves.\n", + "\n", + "| model | weights from | licence |\n", + "|---|---|---|\n", + "| `openbind0` | [OpenBind0 / OpenFold3 v0.5.0](https://github.com/aqlaboratory/openfold-3/releases/tag/v0.5.0) (AlQuraishi Lab) | Apache-2.0 |\n", + "| `openfold3` | [OpenFold3 preview-2](https://github.com/aqlaboratory/openfold) (AlQuraishi Lab) | Apache-2.0 |\n", + "| `boltz2` | [Boltz-2](https://github.com/jwohlwend/boltz) (MIT / Jeremy Wohlwend et al.) | MIT |\n", + "| `protenix2` | [Protenix-v2](https://github.com/bytedance/Protenix) (ByteDance) | Apache-2.0 |\n", + "| `rosettafold3` | [RoseTTAFold3](https://github.com/RosettaCommons/foundry) (RosettaCommons) | BSD-3-Clause |\n", + "| `chai1` | [chai-1](https://github.com/chaidiscovery/chai-lab) (Chai Discovery) | Apache-2.0 |\n", + "| `intellifold2` | [IntelliFold-v2](https://huggingface.co/intelligenAI/intellifold) (IntelligenAI) | Apache-2.0 |\n", + "| `opendde` | [OpenDDE](https://huggingface.co/aurekaresearch/OpenDDE) (Aureka Research) | Apache-2.0 |\n", + "| `esmfold2` | [ESMFold2](https://huggingface.co/biohub/ESMFold2) (Arc Institute / Biohub) | MIT |\n", + "| `esmfold2_lm600m` | ESMFold2 against the 600M ESM-C tower | MIT |\n", + "| `esmfold2_lm300m` | ESMFold2 against the 300M ESM-C tower | MIT |\n", + "| `af2_ptm` | AlphaFold 2 monomer pTM (DeepMind) | CC BY 4.0 |\n", + "| `af2_multimer` | AlphaFold 2 multimer v3 (DeepMind) | CC BY 4.0 |\n", + "| `alphafold3` | Google DeepMind's own parameters | [AF3 terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md) |\n", + "\n", + "Twelve of them run through the **same** JAX/Haiku AlphaFold 3 graph — only the weights and a few gated forward branches differ — so the input box, the MSA path, the outputs, the confidence metrics and the plots are identical whichever you pick. Switch models by changing one dropdown and re-running.\n", + "\n", + "MSA generation via the [ColabFold](https://github.com/sokrypton/ColabFold) MMseqs2 server — **no local databases required**. Attention/XLA flags are chosen automatically for your runtime (T4, L4/Ada, A100/H100, 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", + "- Mirdita et al. (2022) ColabFold. *Nature Methods* [doi:10.1038/s41592-022-01488-1](https://doi.org/10.1038/s41592-022-01488-1)\n", + "- Whichever model you run — please cite it too; each links to its source above.\n", + "\n", + "**Credits:** AF3 code: Google DeepMind (Apache 2.0) · weights: each model's authors, as listed.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "install", + "metadata": { + "cellView": "form", + "id": "install" + }, + "outputs": [], + "source": [ + "#@title Install dependencies (~3 mins)\n", + "%%time\n", + "import os, time, glob, shutil, sys\n", + "\n", + "model = \"openbind0\" #@param [\"openbind0\", \"openfold3\", \"boltz2\", \"protenix2\", \"rosettafold3\", \"chai1\", \"intellifold2\", \"opendde\", \"esmfold2\", \"esmfold2_lm600m\", \"esmfold2_lm300m\", \"alphafold3\", \"af2_ptm\", \"af2_multimer\"]\n", + "#@markdown - **model**: which set of weights to run. All of them use the same AlphaFold 3\n", + "#@markdown graph, so everything downstream is identical. `openbind0` is OpenFold3's current\n", + "#@markdown release and a good default; `openfold3` is their earlier preview-2, kept because\n", + "#@markdown earlier results used it. The three `esmfold2*` entries fold from ESM-C instead\n", + "#@markdown of an MSA -- single sequence, no search -- and differ only in the size of that\n", + "#@markdown language model (6B, 600M, 300M). `chai1` and `esmfold2*` download and run their\n", + "#@markdown language model automatically. `alphafold3` fetches Google DeepMind's own\n", + "#@markdown parameters and is subject to the AF3 terms of use.\n", + "\n", + "persist_cache_to_drive = False #@param {type:\"boolean\"}\n", + "#@markdown - **persist_cache_to_drive**: keep the compiled model in your Google Drive so\n", + "#@markdown the next session does not recompile. Measured on a 68-residue input: the first\n", + "#@markdown prediction takes **69 s** with a cold cache and **16 s** with a warm one, so this\n", + "#@markdown is worth about **53 s per session** (more for longer inputs). Colab wipes `/tmp`\n", + "#@markdown between sessions, which is why it has to go somewhere else to survive. Leaving\n", + "#@markdown it off costs only that recompile; it never changes a result.\n", + "\n", + "# PINNED, both halves. Until 2026-09-16 this installed the v3.1.5 wheel for its\n", + "# compiled extension and then overlaid the Python half from the BRANCH HEAD --\n", + "# so the notebook mixed a fixed binary with a moving source tree, and two runs\n", + "# on different days could be different code. 3.1.7 is published on PyPI\n", + "# (`alphafold3-colabfold`, cp312/cp313/cp314 manylinux + macOS arm64) and its\n", + "# Python half already knows every model, so the overlay is gone and both the\n", + "# package and run_alphafold.py come from one tag.\n", + "VERSION = '3.1.7'\n", + "NATIVE_DIR = 'af3_native_weights'\n", + "AF3_WEIGHTS_URL = 'https://storage.googleapis.com/alphafold3/af3.bin.zst'\n", + "IS_AF3 = (model == 'alphafold3')\n", + "# AlphaFold 2 is a SIBLING NETWORK, not one of the AF3-family ports: MSA row and\n", + "# column attention into an IPA head, reached through the same CLI and writing the\n", + "# same outputs. Its parameters are DeepMind's own release under CC BY 4.0, so they\n", + "# are fetched from source. Protein only -- a ligand or nucleotide in the input\n", + "# raises rather than folding the protein part and saying nothing.\n", + "IS_AF2 = model.startswith('af2_')\n", + "AF2_DIR = 'af2_params'\n", + "# int8 everywhere: same weights stored 8-bit and expanded on load, which is\n", + "# what keeps a Colab download to a few hundred MB. Not a knob -- there is no\n", + "# reason to pick anything else here, and AlphaFold 3's own parameters come\n", + "# from Google as float32 regardless.\n", + "PRECISION = 'fp32' if (IS_AF3 or IS_AF2) else 'int8'\n", + "\n", + "if not os.path.isfile('ALPHAFOLD3_READY'):\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", + " # THE FAT WHEEL, from the GitHub release -- not the slim one on PyPI.\n", + " # `alphafold3.cpp` needs libcifpp's components.cif (518 MB raw, 120 MB\n", + " # zipped) and cannot import without it:\n", + " # ImportError: Could not find the libcifpp components.cif file.\n", + " # With the data the wheel is 130 MB, over PyPI's 100 MB per-file limit, so\n", + " # PyPI carries the slim build (correct for anyone who provisions the data\n", + " # themselves) and the release carries `+data`, which is self-contained.\n", + " _whl = (f'https://github.com/sokrypton/alphafold3/releases/download/v{VERSION}'\n", + " f'/alphafold3_colabfold-{VERSION}%2Bdata-cp313-cp313'\n", + " f'-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl')\n", + " os.system(f\"pip install -q --no-deps '{_whl}'\")\n", + " # `run_alphafold.py` is a top-level script, not part of the package\n", + " # (`wheel.packages = [\"src/alphafold3\"]`), so the wheel does not carry it.\n", + " # Fetch it AT THE TAG so the driver and the library are the same commit.\n", + " os.system(f'wget -q -O run_alphafold.py https://raw.githubusercontent.com'\n", + " f'/sokrypton/alphafold3/v{VERSION}/run_alphafold.py')\n", + " # haiku 0.0.17 still calls the moved `jax.core.DropVar`; checked against the\n", + " # installed 0.0.17 tree, this one is still needed. (A second sed for\n", + " # `jax.core.get_opaque_trace_state` used to sit here and never matched --\n", + " # base.py reaches it through a `jax_core` alias and already falls back to\n", + " # `jex_core` itself, so it was only ever a no-op.)\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", + " os.system('touch ALPHAFOLD3_READY')\n", + " print('Packages installed.')\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", + "# need more shared memory than Ada cards have -> 'Shared memory size limit exceeded' at\n", + "# launch (which its trace-time fallback can't catch). Restrict Triton to true datacenter\n", + "# GPUs (A100 cc 8.0, H100 cc 9.0+); everything else uses XLA, exactly like the T4 path.\n", + "try:\n", + " import tokamax\n", + " _gu = os.path.join(os.path.dirname(tokamax.__file__), '_src', 'gpu_utils.py')\n", + " _s = open(_gu).read()\n", + " _old = 'return float(device.compute_capability) >= 8.0'\n", + " _new = ('cc = float(device.compute_capability)\\n'\n", + " ' return cc == 8.0 or cc >= 9.0 # datacenter only; Ada/L4 (8.6/8.9) lack shared memory')\n", + " if _old in _s:\n", + " open(_gu, 'w').write(_s.replace(_old, _new))\n", + " print('Patched tokamax: Triton restricted to datacenter GPUs (L4/Ada -> XLA).')\n", + "except Exception as _e:\n", + " print(f'(tokamax patch skipped: {_e})')\n", + "\n", + "# Weights, in the background. The ported models are fetched by the same code the run\n", + "# uses (alphafold3.model.weights.ensure_weights), so the run finds them already there\n", + "# and the cache layout cannot drift between the two. AlphaFold 3's own parameters are\n", + "# not ours to redistribute, so those come straight from Google.\n", + "STAMP = f'WEIGHTS_DONE_{model}_{PRECISION}'\n", + "if not os.path.isfile(STAMP):\n", + " if IS_AF2:\n", + " print('Downloading official AlphaFold 2 parameters (CC BY 4.0)...')\n", + " with open('prefetch_af2.py', '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')\n", + " os.system(f'(python prefetch_af2.py {AF2_DIR} && touch {STAMP}) &')\n", + " elif IS_AF3:\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 {STAMP}) &')\n", + " else:\n", + " print(f'Downloading {model} weights...')\n", + " with open('prefetch_weights.py', 'w') as fh:\n", + " fh.write('import sys\\n'\n", + " 'from alphafold3.model import weights\\n'\n", + " 'print(weights.ensure_weights(sys.argv[1], None, precision=sys.argv[2]))\\n')\n", + " os.system(f'(python prefetch_weights.py {model} {PRECISION} && touch {STAMP}) &')\n", + "\n", + "# Where the compiled model is cached. /tmp is wiped when the VM goes away, so a\n", + "# fresh session recompiles (~53 s on a small input); Drive survives. Opt-in, and\n", + "# the run falls back to /tmp if the mount does not work rather than failing.\n", + "CACHE_DIR = '/tmp/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", + "# 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", + "for sentinel in (STAMP, 'DATA_DONE'):\n", + " while not os.path.isfile(sentinel):\n", + " time.sleep(5)\n", + " print(f'{sentinel} ✓')\n", + "\n", + "if IS_AF3 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", + "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" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "input", + "metadata": { + "cellView": "form", + "id": "input" + }, + "outputs": [], + "source": [ + "#@title Input sequences\n", + "import re, os, json, hashlib\n", + "\n", + "#@markdown ### Molecules\n", + "#@markdown Separate multiple chains within a box using `:` (extra colons are fine: `A::::B` == `A:B`). Leave a box empty if unused; full details in the Instructions cell.\n", + "protein = 'PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK' #@param {type:\"string\"}\n", + "dna = '' #@param {type:\"string\"}\n", + "rna = '' #@param {type:\"string\"}\n", + "ligand_ccd = '' #@param {type:\"string\"}\n", + "ligand_smiles = '' #@param {type:\"string\"}\n", + "\n", + "#@markdown ### Run settings\n", + "jobname = 'test' #@param {type:\"string\"}\n", + "msa_mode = \"mmseqs2_server\" #@param [\"mmseqs2_server\", \"single_sequence\"]\n", + "seeds = '1' #@param {type:\"string\"}\n", + "on_existing = \"overwrite\" #@param [\"overwrite\", \"skip\"]\n", + "#@markdown - `msa_mode`: `single_sequence` skips the MSA (faster, lower accuracy).\n", + "#@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", + "# 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", + " return [e for e in (''.join(tok.split()) for tok in s.split(':')) if e]\n", + "\n", + "prot_seqs = [e.upper() for e in split_entries(protein)]\n", + "dna_seqs = [e.upper() for e in split_entries(dna)]\n", + "rna_seqs = [e.upper() for e in split_entries(rna)]\n", + "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", + "# Build AF3 chain entities (IDs A, B, C, ... in canonical order)\n", + "CHAIN_IDS = list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz')\n", + "chains, prot_groups, idx = [], {}, 0\n", + "\n", + "for seq in prot_seqs:\n", + " cid = CHAIN_IDS[idx]; idx += 1\n", + " if seq in prot_groups: # merge identical seqs -> homo-oligomer\n", + " ent = prot_groups[seq]\n", + " ids = ent['id'] if isinstance(ent['id'], list) else [ent['id']]\n", + " ent['id'] = ids + [cid]\n", + " else:\n", + " ent = {'id': cid, 'sequence': seq, 'templates': []}\n", + " if msa_mode == 'single_sequence':\n", + " ent.update({'unpairedMsa': f'>query\\n{seq}\\n', 'pairedMsa': ''})\n", + " prot_groups[seq] = ent\n", + " chains.append({'protein': ent})\n", + "\n", + "for seq in rna_seqs:\n", + " c = {'id': CHAIN_IDS[idx], 'sequence': seq}\n", + " if msa_mode == 'single_sequence':\n", + " c['unpairedMsa'] = f'>query\\n{seq}\\n'\n", + " chains.append({'rna': c}); idx += 1\n", + "\n", + "for seq in dna_seqs:\n", + " chains.append({'dna': {'id': CHAIN_IDS[idx], 'sequence': seq}}); idx += 1\n", + "\n", + "for code in ccd_codes:\n", + " chains.append({'ligand': {'id': CHAIN_IDS[idx], 'ccdCodes': [code]}}); idx += 1\n", + "\n", + "for smiles in smiles_strs:\n", + " chains.append({'ligand': {'id': CHAIN_IDS[idx], 'smiles': smiles}}); idx += 1\n", + "\n", + "if not chains:\n", + " raise ValueError('No valid input found - fill in at least one box.')\n", + "\n", + "# Seeds: pull out integers regardless of separators, dedupe, default to [1]\n", + "seed_list = []\n", + "for tok in re.findall(r'\\d+', seeds):\n", + " v = int(tok)\n", + " if v not in seed_list:\n", + " seed_list.append(v)\n", + "if not seed_list:\n", + " seed_list = [1]\n", + "\n", + "# Deterministic, lower-cased job name from inputs+seeds.\n", + "# Same input+seeds -> same folder (so re-runs reuse it instead of piling up).\n", + "# Lower-cased to match run_alphafold.py's sanitised_name() output directory.\n", + "def _flat(mol):\n", + " if 'sequence' in mol: return mol['sequence']\n", + " if 'ccdCodes' in mol: return ','.join(mol['ccdCodes'])\n", + " return mol.get('smiles', '?')\n", + "flat = ':'.join(_flat(list(c.values())[0]) for c in chains) + '|seeds=' + ','.join(map(str, seed_list))\n", + "basejob = (re.sub(r'\\W+', '', ''.join(jobname.split())) or 'job').lower()\n", + "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", + "OUTPUT_DIR = 'af3_output'\n", + "job_dir = f'{OUTPUT_DIR}/{jobname}'\n", + "\n", + "fold_input = {\n", + " 'name': jobname,\n", + " 'sequences': chains,\n", + " 'modelSeeds': seed_list,\n", + " 'dialect': 'alphafold3',\n", + " 'version': 1,\n", + "}\n", + "os.makedirs(INPUT_DIR, exist_ok=True)\n", + "json_path = f'{INPUT_DIR}/{jobname}.json'\n", + "with open(json_path, 'w') as f:\n", + " json.dump(fold_input, f, indent=2)\n", + "\n", + "print(f'Job \"{jobname}\" -> results will be written to {job_dir}/')\n", + "fold_input\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "run", + "metadata": { + "cellView": "form", + "id": "run" + }, + "outputs": [], + "source": [ + "#@title Run the model\n", + "%%time\n", + "import os, shutil, subprocess, glob\n", + "\n", + "#@markdown Inference settings (defaults match AlphaFold 3 - increase 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", + "\n", + "num_recycles = max(1, int(num_recycles))\n", + "num_diffusion_samples = max(1, int(num_diffusion_samples))\n", + "\n", + "os.makedirs(OUTPUT_DIR, exist_ok=True)\n", + "\n", + "# Re-run policy (one folder per job, no timestamped duplicates):\n", + "# overwrite -> wipe this job's folder and recompute\n", + "# skip -> if a finished result (.cif) is already there, don't recompute\n", + "have_results = os.path.isdir(job_dir) and any(f.endswith('.cif') for f in os.listdir(job_dir))\n", + "run_it = not (on_existing == 'skip' and have_results)\n", + "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", + "def detect_device():\n", + " try:\n", + " out = subprocess.run(\n", + " ['nvidia-smi', '--query-gpu=compute_cap', '--format=csv,noheader'],\n", + " capture_output=True, text=True, timeout=15)\n", + " caps = [float(x) for x in out.stdout.split() if x.strip()]\n", + " if caps:\n", + " return 'gpu', min(caps)\n", + " except Exception:\n", + " pass\n", + " return 'cpu', None\n", + "\n", + "device, cap = detect_device()\n", + "nojit = False\n", + "xla_flags = [] # extra XLA flags to export for this device (per AlphaFold 3's guidance)\n", + "\n", + "if device == 'cpu':\n", + " flash_impl = 'xla'\n", + " nojit = True\n", + " print('No GPU detected - running on CPU with XLA attention + --nojit (slow, but avoids the compile).')\n", + "elif cap < 8.0:\n", + " # T4 / V100 (cc 7.x): XLA attention; disable the custom-kernel fusion pass.\n", + " # (Triton GEMM is not supported on these cards, so it is not disabled here.)\n", + " flash_impl = 'xla'\n", + " xla_flags = ['--xla_disable_hlo_passes=custom-kernel-fusion-rewriter']\n", + " print(f'Pre-Ampere GPU (compute capability {cap}) - XLA attention + custom-kernel fusion disabled.')\n", + "elif 8.0 < cap < 9.0:\n", + " # L4 / Ada / consumer Ampere (cc 8.6 / 8.9): limited shared memory. XLA's Triton GEMM\n", + " # kernels exceed it ('Shared memory size limit exceeded'), so disable Triton GEMM\n", + " # (falls back to cuBLAS) and use XLA attention to also avoid the Triton attention kernel.\n", + " flash_impl = 'xla'\n", + " xla_flags = ['--xla_gpu_enable_triton_gemm=false']\n", + " print(f'Ada/consumer GPU (compute capability {cap}) - XLA attention + Triton GEMM disabled (shared-memory limit).')\n", + "else:\n", + " # A100 (cc 8.0) and H100 (cc 9.0+): ample shared memory. Triton flash attention,\n", + " # with Triton GEMM disabled per AlphaFold 3's recommended XLA_FLAGS.\n", + " flash_impl = 'triton'\n", + " xla_flags = ['--xla_gpu_enable_triton_gemm=false']\n", + " print(f'Datacenter GPU (compute capability {cap}) - Triton flash attention + Triton GEMM disabled.')\n", + "\n", + "# Export XLA flags so the child shell (and JAX inside it) inherit them.\n", + "cur = os.environ.get('XLA_FLAGS', '')\n", + "for f in xla_flags:\n", + " if f not in cur:\n", + " cur = (cur + ' ' + f).strip()\n", + "if cur:\n", + " os.environ['XLA_FLAGS'] = cur\n", + "\n", + "print('XLA_FLAGS =', os.environ.get('XLA_FLAGS', '(unset)'))\n", + "\n", + "# Weights. Every ported model resolves its own cache directory (populated by the\n", + "# install cell), so --model_dir is passed only for the two whose parameters come\n", + "# from DeepMind directly: AlphaFold 3's, and AlphaFold 2's (CC BY 4.0, fetched\n", + "# into af2_params by the install cell).\n", + "print(f'Model: {model}')\n", + "\n", + "cmd = [\n", + " 'python', 'run_alphafold.py',\n", + " f'--json_path={json_path}',\n", + " f'--model={model}',\n", + " '--norun_data_pipeline',\n", + " f'--output_dir={OUTPUT_DIR}',\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", + " f'--num_diffusion_samples={num_diffusion_samples}',\n", + "]\n", + "if msa_mode == 'mmseqs2_server':\n", + " cmd.append('--use_msa_server')\n", + "# chai-1 folds from ESM2 and ESMFold2 from ESM-C; without it they are a\n", + "# different model, not a slightly worse one (a natural protein goes to 5.70 A\n", + "# where chai-1 reaches 0.642, and an ESMFold2 variant with no MSA encoder has\n", + "# nothing left to fold from at all). Both towers run in-process and download on\n", + "# demand, which is why run_alphafold makes it opt-in and this passes it.\n", + "if model == 'chai1' or model.startswith('esmfold2'):\n", + " cmd.append('--use_esm_embeddings')\n", + "if nojit:\n", + " cmd.append('--nojit')\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", + "if run_it:\n", + " print(cmd)\n", + " !{cmd}\n", + " print(f'\\nDone -> {job_dir}/')\n", + "else:\n", + " print(f'Skipping: results already exist in {job_dir}/ (set on_existing=overwrite to recompute).')\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "display3d", + "metadata": { + "cellView": "form", + "id": "display3d" + }, + "outputs": [], + "source": [ + "#@title Display structures + PAE (py2Dmol)\n", + "import csv, glob, os, json\n", + "import numpy as np\n", + "import py2Dmol\n", + "\n", + "load_as_frames = True #@param {type:\"boolean\"}\n", + "viewer_size = 400\n", + "#@markdown All predicted models load together, best first (**rank_1, rank_2, ...**), each with its own PAE.\n", + "#@markdown - `load_as_frames` **off** -> pick a model from the dropdown.\n", + "#@markdown - `load_as_frames` **on** -> models become frames you can play through (press play / drag the slider).\n", + "#@markdown - The interactive PAE matrix sits beside the structure; click or drag-box on it to highlight residues.\n", + "\n", + "# All models in rank order (best first): from the ranking CSV, fall back to globbing.\n", + "def collect_models():\n", + " ranking_csv = f'{job_dir}/{jobname}_ranking_scores.csv'\n", + " cifs = []\n", + " if os.path.exists(ranking_csv):\n", + " rows = []\n", + " with open(ranking_csv) as f:\n", + " for r in csv.DictReader(f):\n", + " rows.append((float(r['ranking_score']), int(r['seed']), int(r['sample'])))\n", + " for _, seed, sample in sorted(rows, reverse=True):\n", + " d = f'{job_dir}/seed-{seed}_sample-{sample}'\n", + " hit = sorted(glob.glob(f'{d}/*_model.cif')) or sorted(glob.glob(f'{d}/*.cif'))\n", + " if hit:\n", + " cifs.append(hit[0])\n", + " if not cifs:\n", + " cifs = (sorted(glob.glob(f'{job_dir}/**/*_model.cif', recursive=True))\n", + " or sorted(glob.glob(f'{job_dir}/**/*.cif', recursive=True)))\n", + " return cifs\n", + "\n", + "# Per-model PAE: confidences.json next to the CIF, else the top-level one.\n", + "def load_pae(cif):\n", + " d = os.path.dirname(cif)\n", + " cands = [p for p in glob.glob(f'{d}/*_confidences.json')\n", + " if 'summary' not in os.path.basename(p)]\n", + " if not cands:\n", + " top = f'{job_dir}/{jobname}_confidences.json'\n", + " cands = [top] if os.path.exists(top) else []\n", + " if cands:\n", + " pae = json.load(open(cands[0])).get('pae')\n", + " if pae is not None:\n", + " return np.asarray(pae, dtype=float)\n", + " return None\n", + "\n", + "cifs = collect_models()\n", + "if not cifs:\n", + " raise FileNotFoundError(f'No model CIFs found in {job_dir}/')\n", + "print(f'Loaded {len(cifs)} model(s) from {job_dir}/'\n", + " + (' (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", + "for i, cif in enumerate(cifs, start=1):\n", + " pae = load_pae(cif)\n", + " if load_as_frames:\n", + " viewer.add_pdb(cif, name='models', paes=pae) # same name -> frames (play through)\n", + " else:\n", + " viewer.add_pdb(cif, name=f'rank_{i}', paes=pae) # distinct names -> dropdown of objects\n", + "viewer.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "plots", + "metadata": { + "cellView": "form", + "id": "plots" + }, + "outputs": [], + "source": [ + "#@title Quality metrics and plots\n", + "import json, os\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "conf_path = f'{OUTPUT_DIR}/{jobname}/{jobname}_confidences.json'\n", + "summ_path = f'{OUTPUT_DIR}/{jobname}/{jobname}_summary_confidences.json'\n", + "\n", + "with open(conf_path) as f:\n", + " conf = json.load(f)\n", + "with open(summ_path) as f:\n", + " summ = json.load(f)\n", + "\n", + "plddts = np.array(conf.get('atom_plddts', conf.get('token_plddts', [])), dtype=float)\n", + "plddt_chain_ids = conf.get('atom_chain_ids', conf.get('token_chain_ids', [])) # pLDDT is per-ATOM\n", + "token_chain_ids = conf.get('token_chain_ids', []) # PAE is per-TOKEN\n", + "pae = np.array(conf.get('pae', []), dtype=float)\n", + "\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", + "mean_plddt = summ.get('mean_plddt')\n", + "if mean_plddt is None and plddts.size:\n", + " mean_plddt = float(np.mean(plddts))\n", + "iptm = summ.get('iptm')\n", + "\n", + "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 — no interface)' if iptm is None else ''))\n", + "print(f'Ranking score : {fmt(summ.get(\"ranking_score\"))}')\n", + "print('=' * 38)\n", + "\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 — 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", + "ax.set_xlim(0, xmax)\n", + "ax.set_ylim(0, 100)\n", + "\n", + "unique_chains = list(dict.fromkeys(plddt_chain_ids))\n", + "if len(unique_chains) > 1:\n", + " colors = plt.cm.tab10(np.linspace(0, 1, len(unique_chains)))\n", + " tcid = np.array(plddt_chain_ids)\n", + " for ch, col in zip(unique_chains, colors):\n", + " y = np.where(tcid == ch, plddts, np.nan) # NaN gaps keep chains as separate lines\n", + " ax.plot(x, y, lw=1.5, color=col, label=f'Chain {ch}')\n", + " for b in [i for i in range(1, len(plddt_chain_ids)) if plddt_chain_ids[i] != plddt_chain_ids[i-1]]:\n", + " ax.axvline(b - 0.5, color='grey', lw=0.6, alpha=0.5)\n", + " ax.legend(loc='lower right', fontsize=8)\n", + "else:\n", + " ax.plot(x, plddts, lw=1.5, color='#1f77b4')\n", + "\n", + "for y in (50, 70, 90):\n", + " ax.axhline(y, ls='--', lw=0.7, color='grey', alpha=0.5)\n", + " ax.text(xmax, y, f' {y}', va='center', ha='left', fontsize=7, color='grey')\n", + "ax.set_xlabel('Atom')\n", + "ax.set_ylabel('pLDDT')\n", + "ax.set_title('Predicted pLDDT per atom')\n", + "\n", + "# PAE matrix\n", + "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 (Å)')\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", + " ax.axvline(b - 0.5, c='black', lw=0.8)\n", + " ax.set_xlabel('Scored residue')\n", + " ax.set_ylabel('Aligned residue')\n", + " ax.set_title('Predicted Aligned Error (PAE)')\n", + "\n", + "plt.tight_layout()\n", + "plt.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "download", + "metadata": { + "cellView": "form", + "id": "download" + }, + "outputs": [], + "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" + ] + }, + { + "cell_type": "markdown", + "id": "instructions", + "metadata": { + "id": "instructions" + }, + "source": [ + "# Instructions \n", + "\n", + "**Quick start**\n", + "1. Pick a **model** in the install cell.\n", + "2. Fill in the sequence(s) in the **Input sequence(s)** cell.\n", + "3. Press **Runtime → Run all**.\n", + "4. The install cell (first run only) downloads the weights for the model you picked and builds AF3 data files in the background — subsequent runs reuse them.\n", + "\n", + "---\n", + "\n", + "## Choosing a model\n", + "\n", + "Pick one in the **model** dropdown of the install cell. Twelve of the fourteen are\n", + "the same AlphaFold 3 network with different trained weights, so nothing else in the\n", + "notebook changes — same input boxes, same MSA path, same outputs and plots. The two\n", + "`af2_*` entries are AlphaFold 2, a different network reached through the same CLI.\n", + "\n", + "| model | notes |\n", + "|---|---|\n", + "| `openbind0` | OpenFold3 v0.5.0 \"OpenBind\", Apache-2.0. The current release, and the default here. |\n", + "| `openfold3` | The earlier OpenFold3 preview-2, Apache-2.0. Kept because earlier results used it. |\n", + "| `boltz2` | MIT. Keeps a modified residue as one token; strong on ligands. |\n", + "| `protenix2` | Apache-2.0. The widest trunk here (pair channel 256), so the slowest. |\n", + "| `rosettafold3` | BSD-3-Clause. Carries chirality features; handles D-amino acids. |\n", + "| `chai1` | Apache-2.0. Folds from ESM2 3B, fetched and run automatically. |\n", + "| `esmfold2` | MIT. Folds from ESM-C instead of an MSA — single sequence, no search. The 6B tower is a 5.1 GB download. |\n", + "| `esmfold2_lm600m` | MIT. Same model against a 600M tower: 0.5 GB instead of 5.1, and no confidence head. |\n", + "| `esmfold2_lm300m` | MIT. The smallest tier, 0.3 GB. Also no confidence head. |\n", + "| `intellifold2` | Apache-2.0. Widened channels (pair 512), largest download. |\n", + "| `opendde` | Apache-2.0. Runs its diffusion on an expanded structural-token set. |\n", + "| `af2_ptm` | AlphaFold 2 monomer pTM, CC BY 4.0. **Protein only** — a ligand or nucleotide in the input raises rather than quietly folding the protein part. Templates use the model_1/model_2 parameter sets, the only monomer ones trained with them. |\n", + "| `af2_multimer` | AlphaFold 2 multimer v3, CC BY 4.0. Protein only, same as above. |\n", + "| `alphafold3` | Google DeepMind's own parameters, under the [AF3 terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md). Publicly downloadable now — no login or key — and fetched into `af3_native_weights/`. run_alphafold prints a reminder of the terms at startup. |\n", + "\n", + "Weights for the eleven ported models are downloaded on first use from\n", + "[sokrypton/af3-any-model](https://huggingface.co/sokrypton/af3-any-model) into a\n", + "per-model cache, so switching models re-downloads only the new one and switching\n", + "back is instant.\n", + "\n", + "---\n", + "\n", + "## Download size\n", + "\n", + "| model | download |\n", + "|---|---|\n", + "| esmfold2_lm300m | 0.12 GB + a 0.3 GB tower |\n", + "| esmfold2_lm600m | 0.12 GB + a 0.5 GB tower |\n", + "| esmfold2 | 0.17 GB + a 5.1 GB tower |\n", + "| protenix2 | 0.18 GB |\n", + "| chai1 | 0.25 GB + a 2.4 GB tower |\n", + "| openbind0 | 0.25 GB |\n", + "| openfold3 | 0.25 GB |\n", + "| rosettafold3 | 0.27 GB |\n", + "| opendde | 0.33 GB |\n", + "| boltz2 | 0.35 GB |\n", + "| intellifold2 | 0.59 GB |\n", + "| af2_ptm / af2_multimer | 3.5 GB (one tar holds every AlphaFold 2 parameter set) |\n", + "\n", + "`chai1` and the `esmfold2*` models also download a protein language model the\n", + "first time they run: 2.4 GB for chai-1, and 5.1 / 0.5 / 0.3 GB for `esmfold2`,\n", + "`esmfold2_lm600m` and `esmfold2_lm300m`.\n", + "\n", + "---\n", + "\n", + "## Sequence input\n", + "\n", + "Each molecule type has its own box. Within a box, separate multiple chains with `:`.\n", + "\n", + "| Box | What goes in it | Example |\n", + "|---|---|---|\n", + "| **protein** | amino-acid sequence(s) | `MKTAY...` or `SEQ1:SEQ2` |\n", + "| **dna** | DNA sequence(s) | `CGCGAATTCGCG` |\n", + "| **rna** | RNA sequence(s) | `GCGGAUUUA` |\n", + "| **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", + "\n", + "## MSA modes\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", + "\n", + "## Output files (inside the downloaded zip)\n", + "\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", + "| `*_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` | The licence notice for whichever weights you ran. |\n", + "\n", + "## Interpreting confidence scores\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", + "\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 weights are a few hundred MB; the language models for `chai1` and `esmfold2*` are larger.\n", + " The install cell downloads in the background and waits for it automatically.\n", + "- **Switching models re-downloads**: each model has its own cache directory,\n", + " so switching back to one you have already used is instant.\n", + "\n", + "## License\n", + "\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: Apache-2.0 for openfold3, protenix2, chai1,\n", + "intellifold2 and opendde; MIT for boltz2; BSD-3-Clause for rosettafold3. Outputs from\n", + "any of those seven are **not** subject to Google DeepMind's AlphaFold 3 Output Terms of\n", + "Use and may be used freely, including commercially.\n", + "\n", + "`alphafold3` is the exception: DeepMind's parameters carry their own\n", + "[terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md),\n", + "and the outputs carry DeepMind's output terms. Every run writes a `TERMS_OF_USE.md`\n", + "naming the licence that actually applies to it.\n", + "\n", + "## Bugs / feedback\n", + "\n", + "Report issues at https://github.com/sokrypton/colabfold/issues" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [], + "include_colab_link": true + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "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 From 0a7d59db5ef768749954b88d86356d270c84f20b Mon Sep 17 00:00:00 2001 From: Sergey Ovchinnikov Date: Wed, 16 Sep 2026 21:14:13 -0400 Subject: [PATCH 02/11] Created using Colab --- ColabFold2_preview.ipynb | 435 ++++++++++++++++++++++----------------- 1 file changed, 249 insertions(+), 186 deletions(-) diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb index 7e1ab5d61..89d2cd62c 100644 --- a/ColabFold2_preview.ipynb +++ b/ColabFold2_preview.ipynb @@ -19,35 +19,24 @@ "source": [ "# ColabFold2 preview\n", "\n", - "Predict protein, RNA, DNA and small-molecule structures with [AlphaFold 3](https://www.nature.com/articles/s41586-024-07487-w), running **any** of fourteen models through one implementation. Pick a model in the install cell — the weights download themselves.\n", - "\n", - "| model | weights from | licence |\n", - "|---|---|---|\n", - "| `openbind0` | [OpenBind0 / OpenFold3 v0.5.0](https://github.com/aqlaboratory/openfold-3/releases/tag/v0.5.0) (AlQuraishi Lab) | Apache-2.0 |\n", - "| `openfold3` | [OpenFold3 preview-2](https://github.com/aqlaboratory/openfold) (AlQuraishi Lab) | Apache-2.0 |\n", - "| `boltz2` | [Boltz-2](https://github.com/jwohlwend/boltz) (MIT / Jeremy Wohlwend et al.) | MIT |\n", - "| `protenix2` | [Protenix-v2](https://github.com/bytedance/Protenix) (ByteDance) | Apache-2.0 |\n", - "| `rosettafold3` | [RoseTTAFold3](https://github.com/RosettaCommons/foundry) (RosettaCommons) | BSD-3-Clause |\n", - "| `chai1` | [chai-1](https://github.com/chaidiscovery/chai-lab) (Chai Discovery) | Apache-2.0 |\n", - "| `intellifold2` | [IntelliFold-v2](https://huggingface.co/intelligenAI/intellifold) (IntelligenAI) | Apache-2.0 |\n", - "| `opendde` | [OpenDDE](https://huggingface.co/aurekaresearch/OpenDDE) (Aureka Research) | Apache-2.0 |\n", - "| `esmfold2` | [ESMFold2](https://huggingface.co/biohub/ESMFold2) (Arc Institute / Biohub) | MIT |\n", - "| `esmfold2_lm600m` | ESMFold2 against the 600M ESM-C tower | MIT |\n", - "| `esmfold2_lm300m` | ESMFold2 against the 300M ESM-C tower | MIT |\n", - "| `af2_ptm` | AlphaFold 2 monomer pTM (DeepMind) | CC BY 4.0 |\n", - "| `af2_multimer` | AlphaFold 2 multimer v3 (DeepMind) | CC BY 4.0 |\n", - "| `alphafold3` | Google DeepMind's own parameters | [AF3 terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md) |\n", - "\n", - "Twelve of them run through the **same** JAX/Haiku AlphaFold 3 graph — only the weights and a few gated forward branches differ — so the input box, the MSA path, the outputs, the confidence metrics and the plots are identical whichever you pick. Switch models by changing one dropdown and re-running.\n", - "\n", - "MSA generation via the [ColabFold](https://github.com/sokrypton/ColabFold) MMseqs2 server — **no local databases required**. Attention/XLA flags are chosen automatically for your runtime (T4, L4/Ada, A100/H100, 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", - "- Mirdita et al. (2022) ColabFold. *Nature Methods* [doi:10.1038/s41592-022-01488-1](https://doi.org/10.1038/s41592-022-01488-1)\n", - "- Whichever model you run — please cite it too; each links to its source above.\n", - "\n", - "**Credits:** AF3 code: Google DeepMind (Apache 2.0) · weights: each model's authors, as listed.\n" + "Predict protein, RNA, DNA and small-molecule structures with\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", + "databases are needed, and the attention/XLA flags are chosen for whatever GPU you get.\n", + "\n", + "Twelve of the fourteen are the *same* AlphaFold 3 network with different trained\n", + "weights, so switching one dropdown changes nothing else — same inputs, same outputs,\n", + "same plots. The two `af2_*` entries are AlphaFold 2, a different network reached\n", + "through the same CLI.\n", + "\n", + "**Models, licences, download sizes and every option: the Instructions cell at the\n", + "bottom.**\n", + "\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" ] }, { @@ -60,36 +49,42 @@ }, "outputs": [], "source": [ - "#@title Install dependencies (~3 mins)\n", - "%%time\n", + "#@title Install dependencies (~35 s)\n", + "# No `%%time`: a cell magic has to be the FIRST line, and `#@title` already\n", + "# is -- with both, Colab renders the form and every other runner aborts the\n", + "# cell. Timed explicitly below, which also survives being run headlessly.\n", "import os, time, glob, shutil, sys\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", - "#@markdown - **model**: which set of weights to run. All of them use the same AlphaFold 3\n", - "#@markdown graph, so everything downstream is identical. `openbind0` is OpenFold3's current\n", - "#@markdown release and a good default; `openfold3` is their earlier preview-2, kept because\n", - "#@markdown earlier results used it. The three `esmfold2*` entries fold from ESM-C instead\n", - "#@markdown of an MSA -- single sequence, no search -- and differ only in the size of that\n", - "#@markdown language model (6B, 600M, 300M). `chai1` and `esmfold2*` download and run their\n", - "#@markdown language model automatically. `alphafold3` fetches Google DeepMind's own\n", - "#@markdown parameters and is subject to the AF3 terms of use.\n", + "#@markdown - **model**: which weights to run -- everything downstream is identical.\n", + "#@markdown `openbind0` is a good default; the `esmfold2*` entries fold from ESM-C with\n", + "#@markdown no MSA. Licences, download sizes and notes: the Instructions cell.\n", "\n", "persist_cache_to_drive = False #@param {type:\"boolean\"}\n", - "#@markdown - **persist_cache_to_drive**: keep the compiled model in your Google Drive so\n", - "#@markdown the next session does not recompile. Measured on a 68-residue input: the first\n", - "#@markdown prediction takes **69 s** with a cold cache and **16 s** with a warm one, so this\n", - "#@markdown is worth about **53 s per session** (more for longer inputs). Colab wipes `/tmp`\n", - "#@markdown between sessions, which is why it has to go somewhere else to survive. Leaving\n", - "#@markdown it off costs only that recompile; it never changes a result.\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", + "# HEADLESS OVERRIDES. A form field above is a plain assignment, so a notebook\n", + "# run outside Colab -- colab-cli, CI -- has no way to change one. Any field can\n", + "# be set from the environment instead:\n", + "# AF3_NB_OVERRIDES='{\"model\": \"boltz2\", \"msa_mode\": \"single_sequence\"}'\n", + "# Unset, this does nothing at all, which is every interactive run.\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", "# PINNED, both halves. Until 2026-09-16 this installed the v3.1.5 wheel for its\n", "# compiled extension and then overlaid the Python half from the BRANCH HEAD --\n", "# so the notebook mixed a fixed binary with a moving source tree, and two runs\n", - "# on different days could be different code. 3.1.7 is published on PyPI\n", - "# (`alphafold3-colabfold`, cp312/cp313/cp314 manylinux + macOS arm64) and its\n", - "# Python half already knows every model, so the overlay is gone and both the\n", - "# package and run_alphafold.py come from one tag.\n", - "VERSION = '3.1.7'\n", + "# on different days could be different code. `alphafold3-colabfold` is published\n", + "# on PyPI (cp312/cp313/cp314 manylinux + macOS arm64) and its Python half knows\n", + "# every model, so the overlay is gone and both the package and run_alphafold.py\n", + "# come from one tag.\n", + "VERSION = '3.1.9'\n", "NATIVE_DIR = 'af3_native_weights'\n", "AF3_WEIGHTS_URL = 'https://storage.googleapis.com/alphafold3/af3.bin.zst'\n", "IS_AF3 = (model == 'alphafold3')\n", @@ -108,19 +103,55 @@ "\n", "if not os.path.isfile('ALPHAFOLD3_READY'):\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", - " # THE FAT WHEEL, from the GitHub release -- not the slim one on PyPI.\n", - " # `alphafold3.cpp` needs libcifpp's components.cif (518 MB raw, 120 MB\n", - " # zipped) and cannot import without it:\n", + " # THE SLIM WHEEL, from PyPI. It is 9 MB. Until v3.1.8 this had to be the\n", + " # 130 MB `+data` wheel from a GitHub release, because importing\n", + " # `alphafold3.cpp` died with\n", " # ImportError: Could not find the libcifpp components.cif file.\n", - " # With the data the wheel is 130 MB, over PyPI's 100 MB per-file limit, so\n", - " # PyPI carries the slim build (correct for anyone who provisions the data\n", - " # themselves) and the release carries `+data`, which is self-contained.\n", - " _whl = (f'https://github.com/sokrypton/alphafold3/releases/download/v{VERSION}'\n", - " f'/alphafold3_colabfold-{VERSION}%2Bdata-cp313-cp313'\n", - " f'-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl')\n", - " os.system(f\"pip install -q --no-deps '{_whl}'\")\n", + " # unless the 518 MB dictionary was bundled. That error was OURS, not\n", + " # libcifpp's: mkdssp_pybind.cc threw from module REGISTRATION, so the whole\n", + " # extension -- including `cif_dict`, which featurisation needs -- failed to\n", + " # import for the sake of DSSP, which no fold calls. v3.1.8 defers the check\n", + " # to `get_dssp` itself. VERIFIED on the published 3.1.8 wheel in a clean\n", + " # venv with no components.cif anywhere: the extension imports, a\n", + " # protein+ligand input featurises, and get_dssp raises something actionable.\n", + " # ml_collections and dm-tree ARE declared dependencies of the package, but it\n", + " # goes in with --no-deps (so pip does not re-resolve jax and the CUDA stack\n", + " # Colab already has), so every third-party import has to be listed here. Those two are imported ONLY\n", + " # by the af2 path (`af2/model/config.py`, and dm-tree in three more), which is\n", + " # why every af3-family model worked and `--model af2_ptm` died with\n", + " # `ModuleNotFoundError: No module named 'ml_collections'`.\n", + " # The full set under src/alphafold3/af2 is: absl, haiku, jax, ml_collections,\n", + " # numpy, scipy, tree -- the rest are already here or in Colab's base image.\n", + " # MEASURED against a fresh Colab image (2026-09-16, py 3.13.15, T4), not\n", + " # guessed. Already present, so not installed: zstandard 0.25.0, dm-tree\n", + " # 0.1.10, numpy 2.1.3, scipy 1.16.3, absl-py, and 21 nvidia CUDA wheels.\n", + " # Absent, so installed: dm-haiku, rdkit, tokamax, ml_collections, py2Dmol.\n", + " #\n", + " # DROPPED: awscli and py3Dmol were installed and never used -- `aws` is never\n", + " # invoked and only py2Dmol is imported. awscli alone drags in the boto stack.\n", + " # NO JAX PIN. This used to force jax[cuda12]==0.10.1; Colab ships 0.11.1, so\n", + " # the pin downgraded jax AND re-pulled the whole CUDA wheel stack -- the\n", + " # dominant cost of this cell. Measured on a fresh T4 session: these four\n", + " # install in 8.5 s against minutes with the pin, and jax 0.11.1 imports,\n", + " # traces and folds correctly (openbind0, 20 residues, rc=0, 165 atoms).\n", + " #\n", + " # CAVEAT, stated because it is untested rather than dismissed: that check was\n", + " # on a T4, which takes the XLA attention path. tokamax's Triton kernels are\n", + " # restricted to datacenter GPUs below, so an A100/H100 run exercises code a\n", + " # T4 does not. If a datacenter GPU misbehaves, pin jax again here first.\n", + " os.system(\"pip install -q dm-haiku==0.0.17 rdkit==2025.9.4 \\\n", + " tokamax==0.0.11 ml_collections\")\n", + " # py2Dmol from source: the released wheel lags the repo.\n", + " os.system(\"pip install -q git+https://github.com/sokrypton/py2Dmol.git\")\n", + " # aria2c, for AF2 only: its parameter tar is 5.3 GB and a single connection is\n", + " # the bottleneck, not the link (weights._download_parallel uses it when it is\n", + " # on PATH). Every af3-family blob is 130-350 MB, where this would not pay.\n", + " if IS_AF2:\n", + " os.system(\"apt-get -qq install -y aria2 > /dev/null 2>&1\")\n", + " # --no-deps: the package declares jax, and Colab already ships a working\n", + " # one -- resolving its dependencies would re-pull the whole CUDA wheel\n", + " # stack, which is why every third-party import is listed above instead.\n", + " os.system(f'pip install -q --no-deps alphafold3-colabfold=={VERSION}')\n", " # `run_alphafold.py` is a top-level script, not part of the package\n", " # (`wheel.packages = [\"src/alphafold3\"]`), so the wheel does not carry it.\n", " # Fetch it AT THE TAG so the driver and the library are the same commit.\n", @@ -165,20 +196,20 @@ " fh.write('import sys\\n'\n", " 'from alphafold3.model import weights\\n'\n", " 'print(weights.ensure_af2_params(sys.argv[1]))\\n')\n", - " os.system(f'(python prefetch_af2.py {AF2_DIR} && touch {STAMP}) &')\n", + " os.system(f'(python prefetch_af2.py {AF2_DIR} > {STAMP}.log 2>&1 && touch {STAMP}) &')\n", " elif IS_AF3:\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 {STAMP}) &')\n", + " os.system(f'(wget -O {NATIVE_DIR}/af3.bin.zst \"{AF3_WEIGHTS_URL}\" > {STAMP}.log 2>&1 && touch {STAMP}) &')\n", " else:\n", " print(f'Downloading {model} weights...')\n", " with open('prefetch_weights.py', 'w') as fh:\n", " fh.write('import sys\\n'\n", " 'from alphafold3.model import weights\\n'\n", " 'print(weights.ensure_weights(sys.argv[1], None, precision=sys.argv[2]))\\n')\n", - " os.system(f'(python prefetch_weights.py {model} {PRECISION} && touch {STAMP}) &')\n", + " os.system(f'(python prefetch_weights.py {model} {PRECISION} > {STAMP}.log 2>&1 && touch {STAMP}) &')\n", "\n", "# Where the compiled model is cached. /tmp is wiped when the VM goes away, so a\n", "# fresh session recompiles (~53 s on a small input); Drive survives. Opt-in, and\n", @@ -196,13 +227,57 @@ "\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", - "for sentinel in (STAMP, 'DATA_DONE'):\n", + " print('Fetching the CCD components this fold needs...')\n", + " # NOT build_data. That parses libcifpp's whole components.cif -- 51,275\n", + " # components, 518 MB -- into a 505 MB ccd.pickle, and costs 48 s of the\n", + " # session. A fold references about thirty codes.\n", + " #\n", + " # LocalFold's trick: fetch each component from\n", + " # files.rcsb.org/ligands/download/.cif, kilobytes each, and build the\n", + " # two pickles from just those. Measured: 0.6 s for 40 components, a 0.30 MB\n", + " # pickle, and every field byte-identical to libcifpp's for ALA, SER, GOL,\n", + " # ATP, SEP, NAG, DA and U.\n", + " #\n", + " # Both pickles come from the SAME fetched set, so they are self-consistent --\n", + " # NAG lands in GLYCAN_LINKING_LIGANDS exactly as with the full dictionary. A\n", + " # component the input names and we did not fetch raises KeyError, which is\n", + " # loud, rather than being silently mis-bonded; hence the generous code list.\n", + " # `ccd_fetch` ships in the wheel as of v3.1.8. It used to be fetched from\n", + " # `main`, which meant the notebook mixed a tagged package with a moving\n", + " # file -- the exact drift the pinned install above exists to prevent.\n", + " with open('prefetch_ccd.py', 'w') as fh:\n", + " fh.write(\n", + " '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", + " os.system('(python prefetch_ccd.py > DATA_DONE.log 2>&1 && touch DATA_DONE) &')\n", + "\n", + "# A BOUNDED wait. This used to be `while not exists: sleep(5)` with no limit\n", + "# and the background job's output discarded, so a failed download was an\n", + "# indefinite hang with nothing on screen -- which is exactly how it looked for\n", + "# ten minutes on 2026-09-17. Each job now writes .log, and a stall\n", + "# raises with the tail of it rather than waiting for the runtime to be\n", + "# reclaimed. The weights are the slow one: a few hundred MB, so minutes on a\n", + "# poor link, hence 20 of them before giving up.\n", + "def _await(sentinel, limit=1200):\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} ✓ ({time.time() - t0:.0f} s)')\n", + "\n", + "for sentinel in (STAMP, 'DATA_DONE'):\n", + " _await(sentinel)\n", "\n", "if IS_AF3 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", @@ -210,7 +285,8 @@ "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" + " ' features. Expect worse structures than chai-lab itself produces.')\n", + "print(f'Setup took {time.time() - _T0:.0f} s.')\n" ] }, { @@ -243,6 +319,17 @@ "#@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. A form field above is a plain assignment, so a notebook\n", + "# run outside Colab -- colab-cli, CI -- has no way to change one. Any field can\n", + "# be set from the environment instead:\n", + "# AF3_NB_OVERRIDES='{\"model\": \"boltz2\", \"msa_mode\": \"single_sequence\"}'\n", + "# Unset, this does nothing at all, which is every interactive run.\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", @@ -341,14 +428,26 @@ "outputs": [], "source": [ "#@title Run the model\n", - "%%time\n", - "import os, shutil, subprocess, glob\n", + "# No `%%time` -- see the install cell.\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. A form field above is a plain assignment, so a notebook\n", + "# run outside Colab -- colab-cli, CI -- has no way to change one. Any field can\n", + "# be set from the environment instead:\n", + "# AF3_NB_OVERRIDES='{\"model\": \"boltz2\", \"msa_mode\": \"single_sequence\"}'\n", + "# Unset, this does nothing at all, which is every interactive run.\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", @@ -452,8 +551,22 @@ "cmd = ' '.join(cmd)\n", "if run_it:\n", " print(cmd)\n", - " !{cmd}\n", - " print(f'\\nDone -> {job_dir}/')\n", + " # NOT `!{cmd}`. That reports nothing about how the run ended, and the line\n", + " # below used to print `Done -> ...` whatever happened -- so a hard failure\n", + " # (an ImportError, 230 ms) read as a successful fold with no structures.\n", + " # Popen streams the same output AND yields a status.\n", + " _p = subprocess.Popen(cmd, shell=True, 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" ] @@ -652,73 +765,41 @@ "source": [ "# Instructions \n", "\n", - "**Quick start**\n", - "1. Pick a **model** in the install cell.\n", - "2. Fill in the sequence(s) in the **Input sequence(s)** cell.\n", - "3. Press **Runtime → Run all**.\n", - "4. The install cell (first run only) downloads the weights for the model you picked and builds AF3 data files in the background — subsequent runs reuse them.\n", - "\n", - "---\n", - "\n", - "## Choosing a model\n", - "\n", - "Pick one in the **model** dropdown of the install cell. Twelve of the fourteen are\n", - "the same AlphaFold 3 network with different trained weights, so nothing else in the\n", - "notebook changes — same input boxes, same MSA path, same outputs and plots. The two\n", - "`af2_*` entries are AlphaFold 2, a different network reached through the same CLI.\n", - "\n", - "| model | notes |\n", - "|---|---|\n", - "| `openbind0` | OpenFold3 v0.5.0 \"OpenBind\", Apache-2.0. The current release, and the default here. |\n", - "| `openfold3` | The earlier OpenFold3 preview-2, Apache-2.0. Kept because earlier results used it. |\n", - "| `boltz2` | MIT. Keeps a modified residue as one token; strong on ligands. |\n", - "| `protenix2` | Apache-2.0. The widest trunk here (pair channel 256), so the slowest. |\n", - "| `rosettafold3` | BSD-3-Clause. Carries chirality features; handles D-amino acids. |\n", - "| `chai1` | Apache-2.0. Folds from ESM2 3B, fetched and run automatically. |\n", - "| `esmfold2` | MIT. Folds from ESM-C instead of an MSA — single sequence, no search. The 6B tower is a 5.1 GB download. |\n", - "| `esmfold2_lm600m` | MIT. Same model against a 600M tower: 0.5 GB instead of 5.1, and no confidence head. |\n", - "| `esmfold2_lm300m` | MIT. The smallest tier, 0.3 GB. Also no confidence head. |\n", - "| `intellifold2` | Apache-2.0. Widened channels (pair 512), largest download. |\n", - "| `opendde` | Apache-2.0. Runs its diffusion on an expanded structural-token set. |\n", - "| `af2_ptm` | AlphaFold 2 monomer pTM, CC BY 4.0. **Protein only** — a ligand or nucleotide in the input raises rather than quietly folding the protein part. Templates use the model_1/model_2 parameter sets, the only monomer ones trained with them. |\n", - "| `af2_multimer` | AlphaFold 2 multimer v3, CC BY 4.0. Protein only, same as above. |\n", - "| `alphafold3` | Google DeepMind's own parameters, under the [AF3 terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md). Publicly downloadable now — no login or key — and fetched into `af3_native_weights/`. run_alphafold prints a reminder of the terms at startup. |\n", - "\n", - "Weights for the eleven ported models are downloaded on first use from\n", - "[sokrypton/af3-any-model](https://huggingface.co/sokrypton/af3-any-model) into a\n", - "per-model cache, so switching models re-downloads only the new one and switching\n", - "back is instant.\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", - "## Download size\n", - "\n", - "| model | download |\n", - "|---|---|\n", - "| esmfold2_lm300m | 0.12 GB + a 0.3 GB tower |\n", - "| esmfold2_lm600m | 0.12 GB + a 0.5 GB tower |\n", - "| esmfold2 | 0.17 GB + a 5.1 GB tower |\n", - "| protenix2 | 0.18 GB |\n", - "| chai1 | 0.25 GB + a 2.4 GB tower |\n", - "| openbind0 | 0.25 GB |\n", - "| openfold3 | 0.25 GB |\n", - "| rosettafold3 | 0.27 GB |\n", - "| opendde | 0.33 GB |\n", - "| boltz2 | 0.35 GB |\n", - "| intellifold2 | 0.59 GB |\n", - "| af2_ptm / af2_multimer | 3.5 GB (one tar holds every AlphaFold 2 parameter set) |\n", - "\n", - "`chai1` and the `esmfold2*` models also download a protein language model the\n", - "first time they run: 2.4 GB for chai-1, and 5.1 / 0.5 / 0.3 GB for `esmfold2`,\n", - "`esmfold2_lm600m` and `esmfold2_lm300m`.\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", + "| model | weights | licence | download | 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 | 0.25 GB | The current release, and the default here. |\n", + "| `openfold3` | [OpenFold3 preview-2](https://github.com/aqlaboratory/openfold) (AlQuraishi Lab) | Apache-2.0 | 0.25 GB | The earlier preview, kept because earlier results used it. |\n", + "| `boltz2` | [Boltz-2](https://github.com/jwohlwend/boltz) (Wohlwend et al.) | MIT | 0.35 GB | Strong on ligands; keeps a modified residue as one token. |\n", + "| `protenix2` | [Protenix-v2](https://github.com/bytedance/Protenix) (ByteDance) | Apache-2.0 | 0.18 GB | The widest trunk here (pair 256), so the slowest. |\n", + "| `rosettafold3` | [RoseTTAFold3](https://github.com/RosettaCommons/foundry) (RosettaCommons) | BSD-3-Clause | 0.27 GB | Carries chirality features; handles D-amino acids. |\n", + "| `chai1` | [chai-1](https://github.com/chaidiscovery/chai-lab) (Chai Discovery) | Apache-2.0 | 0.25 GB + 2.4 GB tower | Folds from ESM2 3B, fetched and run automatically. |\n", + "| `intellifold2` | [IntelliFold-v2](https://huggingface.co/intelligenAI/intellifold) (IntelligenAI) | Apache-2.0 | 0.59 GB | Widened channels (pair 512), largest ported download. |\n", + "| `opendde` | [OpenDDE](https://huggingface.co/aurekaresearch/OpenDDE) (Aureka Research) | Apache-2.0 | 0.33 GB | Runs its diffusion on an expanded structural-token set. |\n", + "| `esmfold2` | [ESMFold2](https://huggingface.co/biohub/ESMFold2) (Arc Institute / Biohub) | MIT | 0.17 GB + 5.1 GB tower | Folds from ESM-C instead of an MSA — single sequence, no search. |\n", + "| `esmfold2_lm600m` | ESMFold2, 600M tower | MIT | 0.12 GB + 0.5 GB tower | No confidence head. |\n", + "| `esmfold2_lm300m` | ESMFold2, 300M tower | MIT | 0.12 GB + 0.3 GB tower | No confidence head. |\n", + "| `af2_ptm` | AlphaFold 2 monomer pTM (DeepMind) | CC BY 4.0 | 3.5 GB | **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 | (same tar) | 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) | — | Publicly downloadable, no login. run_alphafold prints the terms at startup. |\n", "\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", @@ -726,69 +807,51 @@ "| **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", + "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", - "## Seeds\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", - "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", + "## Output\n", "\n", - "## MSA modes\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", - "\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` | The licence notice for whichever weights you ran. |\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", - "## Interpreting confidence scores\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 weights are a few hundred MB; the language models for `chai1` and `esmfold2*` are larger.\n", - " The install cell downloads in the background and waits for it automatically.\n", - "- **Switching models re-downloads**: each model has its own cache directory,\n", - " so switching back to one you have already used is instant.\n", - "\n", - "## License\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", - "The AlphaFold 3 **source code** is [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0).\n", - "The **weights** are each their own: Apache-2.0 for openfold3, protenix2, chai1,\n", - "intellifold2 and opendde; MIT for boltz2; BSD-3-Clause for rosettafold3. Outputs from\n", - "any of those seven are **not** subject to Google DeepMind's AlphaFold 3 Output Terms of\n", - "Use and may be used freely, including commercially.\n", + "## Licence\n", "\n", - "`alphafold3` is the exception: DeepMind's parameters carry their own\n", - "[terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md),\n", - "and the outputs carry DeepMind's output terms. Every run writes a `TERMS_OF_USE.md`\n", - "naming the licence that actually applies to it.\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/colabfold/issues" + "https://github.com/sokrypton/alphafold3/issues\n" ] } ], From 08cfca998a696cfb0ecaac14a6408da720f19205 Mon Sep 17 00:00:00 2001 From: Sergey Ovchinnikov Date: Wed, 16 Sep 2026 22:07:56 -0400 Subject: [PATCH 03/11] Created using Colab --- ColabFold2_preview.ipynb | 127 +++++++++++++++++++++------------------ 1 file changed, 68 insertions(+), 59 deletions(-) diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb index 89d2cd62c..56fdaabb1 100644 --- a/ColabFold2_preview.ipynb +++ b/ColabFold2_preview.ipynb @@ -26,11 +26,6 @@ "[ColabFold](https://github.com/sokrypton/ColabFold) MMseqs2 server, so no local\n", "databases are needed, and the attention/XLA flags are chosen for whatever GPU you get.\n", "\n", - "Twelve of the fourteen are the *same* AlphaFold 3 network with different trained\n", - "weights, so switching one dropdown changes nothing else — same inputs, same outputs,\n", - "same plots. The two `af2_*` entries are AlphaFold 2, a different network reached\n", - "through the same CLI.\n", - "\n", "**Models, licences, download sizes and every option: the Instructions cell at the\n", "bottom.**\n", "\n", @@ -57,9 +52,6 @@ "_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", - "#@markdown - **model**: which weights to run -- everything downstream is identical.\n", - "#@markdown `openbind0` is a good default; the `esmfold2*` entries fold from ESM-C with\n", - "#@markdown no MSA. Licences, download sizes and notes: the Instructions cell.\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", @@ -84,7 +76,7 @@ "# on PyPI (cp312/cp313/cp314 manylinux + macOS arm64) and its Python half knows\n", "# every model, so the overlay is gone and both the package and run_alphafold.py\n", "# come from one tag.\n", - "VERSION = '3.1.9'\n", + "VERSION = '3.1.10'\n", "NATIVE_DIR = 'af3_native_weights'\n", "AF3_WEIGHTS_URL = 'https://storage.googleapis.com/alphafold3/af3.bin.zst'\n", "IS_AF3 = (model == 'alphafold3')\n", @@ -101,6 +93,17 @@ "# from Google as float32 regardless.\n", "PRECISION = 'fp32' if (IS_AF3 or IS_AF2) else 'int8'\n", "\n", + "\n", + "# Every install below is CHECKED. os.system returns a status and this cell used\n", + "# to throw it away, then print \"Packages installed.\" regardless -- so a pip that\n", + "# found no distribution (PyPI's index lags its file store by a couple of\n", + "# minutes after a release) became `ModuleNotFoundError: No module named\n", + "# 'alphafold3'` from a background job several steps later, with nothing\n", + "# connecting the two.\n", + "def _sh(cmd, what):\n", + " if os.system(cmd) != 0:\n", + " raise RuntimeError(f'{what} failed. The pip output is above.')\n", + "\n", "if not os.path.isfile('ALPHAFOLD3_READY'):\n", " print('Installing packages...')\n", " # THE SLIM WHEEL, from PyPI. It is 9 MB. Until v3.1.8 this had to be the\n", @@ -139,10 +142,11 @@ " # on a T4, which takes the XLA attention path. tokamax's Triton kernels are\n", " # restricted to datacenter GPUs below, so an A100/H100 run exercises code a\n", " # T4 does not. If a datacenter GPU misbehaves, pin jax again here first.\n", - " os.system(\"pip install -q dm-haiku==0.0.17 rdkit==2025.9.4 \\\n", - " tokamax==0.0.11 ml_collections\")\n", + " _sh(\"pip install -q dm-haiku==0.0.17 rdkit==2025.9.4 \"\n", + " \"tokamax==0.0.11 ml_collections\", 'installing dependencies')\n", " # py2Dmol from source: the released wheel lags the repo.\n", - " os.system(\"pip install -q git+https://github.com/sokrypton/py2Dmol.git\")\n", + " _sh(\"pip install -q git+https://github.com/sokrypton/py2Dmol.git\",\n", + " 'installing py2Dmol')\n", " # aria2c, for AF2 only: its parameter tar is 5.3 GB and a single connection is\n", " # the bottleneck, not the link (weights._download_parallel uses it when it is\n", " # on PATH). Every af3-family blob is 130-350 MB, where this would not pay.\n", @@ -151,20 +155,28 @@ " # --no-deps: the package declares jax, and Colab already ships a working\n", " # one -- resolving its dependencies would re-pull the whole CUDA wheel\n", " # stack, which is why every third-party import is listed above instead.\n", - " os.system(f'pip install -q --no-deps alphafold3-colabfold=={VERSION}')\n", + " # Retried: a wheel published minutes ago may not be in the index yet.\n", + " for _try in range(4):\n", + " if os.system(f'pip install -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", " # (`wheel.packages = [\"src/alphafold3\"]`), so the wheel does not carry it.\n", " # Fetch it AT THE TAG so the driver and the library are the same commit.\n", - " os.system(f'wget -q -O run_alphafold.py https://raw.githubusercontent.com'\n", - " f'/sokrypton/alphafold3/v{VERSION}/run_alphafold.py')\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`; checked against the\n", " # installed 0.0.17 tree, this one is still needed. (A second sed for\n", " # `jax.core.get_opaque_trace_state` used to sit here and never matched --\n", " # base.py reaches it through a `jax_core` alias and already falls back to\n", " # `jex_core` itself, so it was only ever a no-op.)\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 # the only proof that any of the above worked\n", " os.system('touch ALPHAFOLD3_READY')\n", - " print('Packages installed.')\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", @@ -225,46 +237,11 @@ " except Exception as _e:\n", " print(f'(Drive mount failed, using {CACHE_DIR}: {_e})')\n", "\n", - "# Build AF3 data files (background, independent of weights)\n", - "if not os.path.isfile('DATA_DONE'):\n", - " print('Fetching the CCD components this fold needs...')\n", - " # NOT build_data. That parses libcifpp's whole components.cif -- 51,275\n", - " # components, 518 MB -- into a 505 MB ccd.pickle, and costs 48 s of the\n", - " # session. A fold references about thirty codes.\n", - " #\n", - " # LocalFold's trick: fetch each component from\n", - " # files.rcsb.org/ligands/download/.cif, kilobytes each, and build the\n", - " # two pickles from just those. Measured: 0.6 s for 40 components, a 0.30 MB\n", - " # pickle, and every field byte-identical to libcifpp's for ALA, SER, GOL,\n", - " # ATP, SEP, NAG, DA and U.\n", - " #\n", - " # Both pickles come from the SAME fetched set, so they are self-consistent --\n", - " # NAG lands in GLYCAN_LINKING_LIGANDS exactly as with the full dictionary. A\n", - " # component the input names and we did not fetch raises KeyError, which is\n", - " # loud, rather than being silently mis-bonded; hence the generous code list.\n", - " # `ccd_fetch` ships in the wheel as of v3.1.8. It used to be fetched from\n", - " # `main`, which meant the notebook mixed a tagged package with a moving\n", - " # file -- the exact drift the pinned install above exists to prevent.\n", - " with open('prefetch_ccd.py', 'w') as fh:\n", - " fh.write(\n", - " '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", - " os.system('(python prefetch_ccd.py > DATA_DONE.log 2>&1 && touch DATA_DONE) &')\n", - "\n", - "# A BOUNDED wait. This used to be `while not exists: sleep(5)` with no limit\n", - "# and the background job's output discarded, so a failed download was an\n", - "# indefinite hang with nothing on screen -- which is exactly how it looked for\n", - "# ten minutes on 2026-09-17. Each job now writes .log, and a stall\n", - "# raises with the tail of it rather than waiting for the runtime to be\n", - "# reclaimed. The weights are the slow one: a few hundred MB, so minutes on a\n", - "# poor link, hence 20 of them before giving up.\n", + "# A BOUNDED wait. This used to be `while not exists: sleep(5)` with the\n", + "# background job's output discarded, so a failed download was an indefinite\n", + "# hang with nothing on screen -- which is exactly how a HuggingFace 429 looked\n", + "# on 2026-09-17. The job writes .log; a stall raises with the tail of\n", + "# it. A few hundred MB, so 20 minutes before giving up.\n", "def _await(sentinel, limit=1200):\n", " t0 = time.time()\n", " while not os.path.isfile(sentinel):\n", @@ -274,10 +251,10 @@ " 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} ✓ ({time.time() - t0:.0f} s)')\n", + " print(f'{sentinel} \\u2713 ({time.time() - t0:.0f} s)')\n", "\n", - "for sentinel in (STAMP, 'DATA_DONE'):\n", - " _await(sentinel)\n", + "\n", + "_await(STAMP)\n", "\n", "if IS_AF3 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", @@ -341,6 +318,38 @@ "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", + "# THE CCD, for the components this input actually names. It has to happen here\n", + "# and not in the install cell: that cell runs BEFORE you type a ligand, so it\n", + "# can only ever fetch the standard residues -- and a fold naming ATP then dies\n", + "# with `ValueError: Unknown residue type ATP`.\n", + "#\n", + "# NOT build_data, which parses libcifpp's whole components.cif (51,275\n", + "# components, 518 MB) into a 505 MB pickle and costs 48 s of the session.\n", + "# LocalFold's trick instead: pull each component from\n", + "# files.rcsb.org/ligands/download/.cif, kilobytes each. Measured at 0.6 s\n", + "# for 40 components, every field byte-identical to libcifpp's for ALA, SER,\n", + "# GOL, ATP, SEP, NAG, DA and U.\n", + "#\n", + "# Both pickles are written from the SAME fetched set, so they stay\n", + "# self-consistent -- NAG lands in GLYCAN_LINKING_LIGANDS exactly as it would\n", + "# with the full dictionary. A code we did not fetch raises KeyError, which is\n", + "# loud, rather than being silently mis-bonded.\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", From 1b312c8dfcb2f7916494975ea7e18bd131f62f80 Mon Sep 17 00:00:00 2001 From: Sergey Ovchinnikov Date: Wed, 16 Sep 2026 22:12:35 -0400 Subject: [PATCH 04/11] fixing style --- ColabFold2_preview.ipynb | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb index 56fdaabb1..41cbbae86 100644 --- a/ColabFold2_preview.ipynb +++ b/ColabFold2_preview.ipynb @@ -642,7 +642,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", From ec208bb5fd017968f3941181f29433a128e4213e Mon Sep 17 00:00:00 2001 From: Sergey Ovchinnikov Date: Wed, 16 Sep 2026 22:41:34 -0400 Subject: [PATCH 05/11] fixing credits --- ColabFold2_preview.ipynb | 33 ++++++++++++++++----------------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb index 41cbbae86..38b648510 100644 --- a/ColabFold2_preview.ipynb +++ b/ColabFold2_preview.ipynb @@ -786,23 +786,22 @@ "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 | download | 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 | 0.25 GB | The current release, and the default here. |\n", - "| `openfold3` | [OpenFold3 preview-2](https://github.com/aqlaboratory/openfold) (AlQuraishi Lab) | Apache-2.0 | 0.25 GB | The earlier preview, kept because earlier results used it. |\n", - "| `boltz2` | [Boltz-2](https://github.com/jwohlwend/boltz) (Wohlwend et al.) | MIT | 0.35 GB | Strong on ligands; keeps a modified residue as one token. |\n", - "| `protenix2` | [Protenix-v2](https://github.com/bytedance/Protenix) (ByteDance) | Apache-2.0 | 0.18 GB | The widest trunk here (pair 256), so the slowest. |\n", - "| `rosettafold3` | [RoseTTAFold3](https://github.com/RosettaCommons/foundry) (RosettaCommons) | BSD-3-Clause | 0.27 GB | Carries chirality features; handles D-amino acids. |\n", - "| `chai1` | [chai-1](https://github.com/chaidiscovery/chai-lab) (Chai Discovery) | Apache-2.0 | 0.25 GB + 2.4 GB tower | Folds from ESM2 3B, fetched and run automatically. |\n", - "| `intellifold2` | [IntelliFold-v2](https://huggingface.co/intelligenAI/intellifold) (IntelligenAI) | Apache-2.0 | 0.59 GB | Widened channels (pair 512), largest ported download. |\n", - "| `opendde` | [OpenDDE](https://huggingface.co/aurekaresearch/OpenDDE) (Aureka Research) | Apache-2.0 | 0.33 GB | Runs its diffusion on an expanded structural-token set. |\n", - "| `esmfold2` | [ESMFold2](https://huggingface.co/biohub/ESMFold2) (Arc Institute / Biohub) | MIT | 0.17 GB + 5.1 GB tower | Folds from ESM-C instead of an MSA — single sequence, no search. |\n", - "| `esmfold2_lm600m` | ESMFold2, 600M tower | MIT | 0.12 GB + 0.5 GB tower | No confidence head. |\n", - "| `esmfold2_lm300m` | ESMFold2, 300M tower | MIT | 0.12 GB + 0.3 GB tower | No confidence head. |\n", - "| `af2_ptm` | AlphaFold 2 monomer pTM (DeepMind) | CC BY 4.0 | 3.5 GB | **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 | (same tar) | 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) | — | Publicly downloadable, no login. run_alphafold prints the terms at startup. |\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 — 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** — 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) | Requires requesting the weights from Google (non-commercial research only, granted at Google's discretion) — not a direct download. run_alphafold prints the terms at startup. |\n", "---\n", "\n", "## Input\n", From 2a0dcd04561f9043d663fb26de764d2361241610 Mon Sep 17 00:00:00 2001 From: Sergey Ovchinnikov Date: Thu, 17 Sep 2026 03:40:54 +0000 Subject: [PATCH 06/11] stop downloading AlphaFold 3's weights, and drop the debugging comments Two changes, both already on sokrypton/alphafold3@main. The model table says DeepMind's parameters are granted on request, for non-commercial research, at Google's discretion -- and the install cell was wgetting them from a public bucket and printing "public, no login required". The URL does serve 1,020 MB, but reachable is not the same as ours to hand out. The `alphafold3` entry now looks for a .bin.zst you placed in af3_native_weights/ yourself and, if there is none, says where to request them. Every other model still fetches its own weights. And most of the comments in the three code cells were a log of last night's debugging -- why %%time was removed, what a HuggingFace 429 looked like, which wheel used to be installed. That belongs in the history, not in front of someone trying to fold a protein. No behaviour change; what is left explains a decision a reader cannot infer. Verified on a cold Colab T4 before pushing: setup 34 s, fold 68 s, ATP ligand present (31 HETATM), zero rasa warnings. Also green on A100 (Triton path) and L4 (XLA fallback). --- ColabFold2_preview.ipynb | 274 +++++++++++++-------------------------- 1 file changed, 89 insertions(+), 185 deletions(-) diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb index 38b648510..4dfb193bc 100644 --- a/ColabFold2_preview.ipynb +++ b/ColabFold2_preview.ipynb @@ -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) — and with thirteen\n", + "[AlphaFold 3](https://www.nature.com/articles/s41586-024-07487-w) \u2014 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", @@ -45,9 +45,7 @@ "outputs": [], "source": [ "#@title Install dependencies (~35 s)\n", - "# No `%%time`: a cell magic has to be the FIRST line, and `#@title` already\n", - "# is -- with both, Colab renders the form and every other runner aborts the\n", - "# cell. Timed explicitly below, which also survives being run headlessly.\n", + "# No `%%time` here: a cell magic must be the first line and `#@title` already is.\n", "import os, time, glob, shutil, sys\n", "_T0 = time.time()\n", "\n", @@ -58,104 +56,42 @@ "#@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", - "# HEADLESS OVERRIDES. A form field above is a plain assignment, so a notebook\n", - "# run outside Colab -- colab-cli, CI -- has no way to change one. Any field can\n", - "# be set from the environment instead:\n", - "# AF3_NB_OVERRIDES='{\"model\": \"boltz2\", \"msa_mode\": \"single_sequence\"}'\n", - "# Unset, this does nothing at all, which is every interactive run.\n", - "import json as _json, os as _os\n", - "for _k, _v in _json.loads(_os.environ.get('AF3_NB_OVERRIDES', '{}')).items():\n", + "# Form fields are plain assignments, so a headless run (colab-cli, CI) sets them\n", + "# through the environment: 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", - "# PINNED, both halves. Until 2026-09-16 this installed the v3.1.5 wheel for its\n", - "# compiled extension and then overlaid the Python half from the BRANCH HEAD --\n", - "# so the notebook mixed a fixed binary with a moving source tree, and two runs\n", - "# on different days could be different code. `alphafold3-colabfold` is published\n", - "# on PyPI (cp312/cp313/cp314 manylinux + macOS arm64) and its Python half knows\n", - "# every model, so the overlay is gone and both the package and run_alphafold.py\n", - "# come from one tag.\n", - "VERSION = '3.1.10'\n", + "VERSION = '3.1.10' # 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", - "# AlphaFold 2 is a SIBLING NETWORK, not one of the AF3-family ports: MSA row and\n", - "# column attention into an IPA head, reached through the same CLI and writing the\n", - "# same outputs. Its parameters are DeepMind's own release under CC BY 4.0, so they\n", - "# are fetched from source. Protein only -- a ligand or nucleotide in the input\n", - "# raises rather than folding the protein part and saying nothing.\n", "IS_AF2 = model.startswith('af2_')\n", - "AF2_DIR = 'af2_params'\n", - "# int8 everywhere: same weights stored 8-bit and expanded on load, which is\n", - "# what keeps a Colab download to a few hundred MB. Not a knob -- there is no\n", - "# reason to pick anything else here, and AlphaFold 3's own parameters come\n", - "# from Google as float32 regardless.\n", + "# int8: the same weights stored 8-bit and expanded on load, which is what keeps\n", + "# the download to a few hundred MB. AF2 and AF3 ship their own float32 files.\n", "PRECISION = 'fp32' if (IS_AF3 or IS_AF2) else 'int8'\n", "\n", "\n", - "# Every install below is CHECKED. os.system returns a status and this cell used\n", - "# to throw it away, then print \"Packages installed.\" regardless -- so a pip that\n", - "# found no distribution (PyPI's index lags its file store by a couple of\n", - "# minutes after a release) became `ModuleNotFoundError: No module named\n", - "# 'alphafold3'` from a background job several steps later, with nothing\n", - "# connecting the two.\n", "def _sh(cmd, what):\n", + " \"\"\"Run a shell command and raise if it fails -- os.system's status is easy to drop.\"\"\"\n", " if os.system(cmd) != 0:\n", - " raise RuntimeError(f'{what} failed. The pip output is above.')\n", + " raise RuntimeError(f'{what} failed. The output is above.')\n", + "\n", "\n", "if not os.path.isfile('ALPHAFOLD3_READY'):\n", " print('Installing packages...')\n", - " # THE SLIM WHEEL, from PyPI. It is 9 MB. Until v3.1.8 this had to be the\n", - " # 130 MB `+data` wheel from a GitHub release, because importing\n", - " # `alphafold3.cpp` died with\n", - " # ImportError: Could not find the libcifpp components.cif file.\n", - " # unless the 518 MB dictionary was bundled. That error was OURS, not\n", - " # libcifpp's: mkdssp_pybind.cc threw from module REGISTRATION, so the whole\n", - " # extension -- including `cif_dict`, which featurisation needs -- failed to\n", - " # import for the sake of DSSP, which no fold calls. v3.1.8 defers the check\n", - " # to `get_dssp` itself. VERIFIED on the published 3.1.8 wheel in a clean\n", - " # venv with no components.cif anywhere: the extension imports, a\n", - " # protein+ligand input featurises, and get_dssp raises something actionable.\n", - " # ml_collections and dm-tree ARE declared dependencies of the package, but it\n", - " # goes in with --no-deps (so pip does not re-resolve jax and the CUDA stack\n", - " # Colab already has), so every third-party import has to be listed here. Those two are imported ONLY\n", - " # by the af2 path (`af2/model/config.py`, and dm-tree in three more), which is\n", - " # why every af3-family model worked and `--model af2_ptm` died with\n", - " # `ModuleNotFoundError: No module named 'ml_collections'`.\n", - " # The full set under src/alphafold3/af2 is: absl, haiku, jax, ml_collections,\n", - " # numpy, scipy, tree -- the rest are already here or in Colab's base image.\n", - " # MEASURED against a fresh Colab image (2026-09-16, py 3.13.15, T4), not\n", - " # guessed. Already present, so not installed: zstandard 0.25.0, dm-tree\n", - " # 0.1.10, numpy 2.1.3, scipy 1.16.3, absl-py, and 21 nvidia CUDA wheels.\n", - " # Absent, so installed: dm-haiku, rdkit, tokamax, ml_collections, py2Dmol.\n", - " #\n", - " # DROPPED: awscli and py3Dmol were installed and never used -- `aws` is never\n", - " # invoked and only py2Dmol is imported. awscli alone drags in the boto stack.\n", - " # NO JAX PIN. This used to force jax[cuda12]==0.10.1; Colab ships 0.11.1, so\n", - " # the pin downgraded jax AND re-pulled the whole CUDA wheel stack -- the\n", - " # dominant cost of this cell. Measured on a fresh T4 session: these four\n", - " # install in 8.5 s against minutes with the pin, and jax 0.11.1 imports,\n", - " # traces and folds correctly (openbind0, 20 residues, rc=0, 165 atoms).\n", - " #\n", - " # CAVEAT, stated because it is untested rather than dismissed: that check was\n", - " # on a T4, which takes the XLA attention path. tokamax's Triton kernels are\n", - " # restricted to datacenter GPUs below, so an A100/H100 run exercises code a\n", - " # T4 does not. If a datacenter GPU misbehaves, pin jax again here first.\n", + " # --no-deps throughout: Colab already ships jax and the CUDA stack, and letting\n", + " # pip re-resolve them re-downloads gigabytes. So every third-party import the\n", + " # package needs is listed here instead.\n", " _sh(\"pip install -q dm-haiku==0.0.17 rdkit==2025.9.4 \"\n", " \"tokamax==0.0.11 ml_collections\", 'installing dependencies')\n", - " # py2Dmol from source: the released wheel lags the repo.\n", - " _sh(\"pip install -q git+https://github.com/sokrypton/py2Dmol.git\",\n", + " _sh(\"pip install -q git+https://github.com/sokrypton/py2Dmol.git\", # wheel lags the repo\n", " 'installing py2Dmol')\n", - " # aria2c, for AF2 only: its parameter tar is 5.3 GB and a single connection is\n", - " # the bottleneck, not the link (weights._download_parallel uses it when it is\n", - " # on PATH). Every af3-family blob is 130-350 MB, where this would not pay.\n", " if IS_AF2:\n", - " os.system(\"apt-get -qq install -y aria2 > /dev/null 2>&1\")\n", - " # --no-deps: the package declares jax, and Colab already ships a working\n", - " # one -- resolving its dependencies would re-pull the whole CUDA wheel\n", - " # stack, which is why every third-party import is listed above instead.\n", - " # Retried: a wheel published minutes ago may not be in the index yet.\n", + " os.system(\"apt-get -qq install -y aria2 > /dev/null 2>&1\") # AF2's tar is 5.3 GB\n", + " # Retried: a wheel published minutes ago may not be in PyPI's index yet.\n", " for _try in range(4):\n", " if os.system(f'pip install -q --no-deps alphafold3-colabfold=={VERSION}') == 0:\n", " break\n", @@ -163,69 +99,64 @@ " 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", - " # (`wheel.packages = [\"src/alphafold3\"]`), so the wheel does not carry it.\n", - " # Fetch it AT THE TAG so the driver and the library are the same commit.\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`; checked against the\n", - " # installed 0.0.17 tree, this one is still needed. (A second sed for\n", - " # `jax.core.get_opaque_trace_state` used to sit here and never matched --\n", - " # base.py reaches it through a `jax_core` alias and already falls back to\n", - " # `jex_core` itself, so it was only ever a no-op.)\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 # the only proof that any of the above worked\n", + " import alphafold3 # the only real proof the install worked\n", " os.system('touch ALPHAFOLD3_READY')\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", - "# need more shared memory than Ada cards have -> 'Shared memory size limit exceeded' at\n", - "# launch (which its trace-time fallback can't catch). Restrict Triton to true datacenter\n", - "# GPUs (A100 cc 8.0, H100 cc 9.0+); everything else uses XLA, exactly like the T4 path.\n", + "# tokamax enables its Triton kernels for every GPU with cc >= 8.0, but they need\n", + "# more shared memory than Ada cards have and fail at launch. Restrict them to\n", + "# datacenter GPUs (A100 cc 8.0, H100 cc 9.0+); Ada/L4 use XLA like a T4 does.\n", "try:\n", " import tokamax\n", " _gu = os.path.join(os.path.dirname(tokamax.__file__), '_src', 'gpu_utils.py')\n", " _s = open(_gu).read()\n", " _old = 'return float(device.compute_capability) >= 8.0'\n", " _new = ('cc = float(device.compute_capability)\\n'\n", - " ' return cc == 8.0 or cc >= 9.0 # datacenter only; Ada/L4 (8.6/8.9) lack shared memory')\n", + " ' return cc == 8.0 or cc >= 9.0 # datacenter only; Ada/L4 lack shared memory')\n", " if _old in _s:\n", " open(_gu, 'w').write(_s.replace(_old, _new))\n", " print('Patched tokamax: Triton restricted to datacenter GPUs (L4/Ada -> XLA).')\n", "except Exception as _e:\n", " print(f'(tokamax patch skipped: {_e})')\n", "\n", - "# Weights, in the background. The ported models are fetched by the same code the run\n", - "# uses (alphafold3.model.weights.ensure_weights), so the run finds them already there\n", - "# and the cache layout cannot drift between the two. AlphaFold 3's own parameters are\n", - "# not ours to redistribute, so those come straight from Google.\n", + "# Weights, fetched in the background by the same code the run uses, so the cache\n", + "# layout cannot drift between the two.\n", "STAMP = f'WEIGHTS_DONE_{model}_{PRECISION}'\n", - "if not os.path.isfile(STAMP):\n", - " if IS_AF2:\n", - " print('Downloading official AlphaFold 2 parameters (CC BY 4.0)...')\n", - " with open('prefetch_af2.py', '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')\n", - " os.system(f'(python prefetch_af2.py {AF2_DIR} > {STAMP}.log 2>&1 && touch {STAMP}) &')\n", - " elif IS_AF3:\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 -O {NATIVE_DIR}/af3.bin.zst \"{AF3_WEIGHTS_URL}\" > {STAMP}.log 2>&1 && touch {STAMP}) &')\n", - " else:\n", - " print(f'Downloading {model} weights...')\n", - " with open('prefetch_weights.py', 'w') as fh:\n", - " fh.write('import sys\\n'\n", - " 'from alphafold3.model import weights\\n'\n", - " 'print(weights.ensure_weights(sys.argv[1], None, precision=sys.argv[2]))\\n')\n", - " os.system(f'(python prefetch_weights.py {model} {PRECISION} > {STAMP}.log 2>&1 && touch {STAMP}) &')\n", - "\n", - "# Where the compiled model is cached. /tmp is wiped when the VM goes away, so a\n", - "# fresh session recompiles (~53 s on a small input); Drive survives. Opt-in, and\n", - "# the run falls back to /tmp if the mount does not work rather than failing.\n", + "if IS_AF3:\n", + " # NOT downloaded. DeepMind's AlphaFold 3 parameters are granted on request,\n", + " # for non-commercial research, at Google's discretion -- they are not ours to\n", + " # fetch on your behalf. Apply at\n", + " # https://docs.google.com/forms/d/e/1FAIpQLSfWZAgo1aYk0O4MuAXZj8xRQ8DafeFJnldNOnh_13qAx2ceZw/viewform\n", + " # and put the file you are given in af3_native_weights/.\n", + " os.makedirs(NATIVE_DIR, exist_ok=True)\n", + " _blobs = glob.glob(f'{NATIVE_DIR}/*.bin.zst')\n", + " if not _blobs:\n", + " raise RuntimeError(\n", + " f'No AlphaFold 3 parameters found in {NATIVE_DIR}/.\\n'\n", + " 'DeepMind grants these on request (non-commercial research); this '\n", + " 'notebook cannot download them for you. Request them at\\n'\n", + " ' https://github.com/google-deepmind/alphafold3#obtaining-model-parameters\\n'\n", + " f'then upload the .bin.zst file into {NATIVE_DIR}/ and re-run this cell.\\n'\n", + " 'Every other model in the dropdown downloads its own weights.')\n", + " print(f'Using your AlphaFold 3 parameters: {_blobs[0]}')\n", + "elif not 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 on a small\n", + "# input); Drive survives. Falls back to /tmp rather than failing.\n", "CACHE_DIR = '/tmp/af3_cache'\n", "if persist_cache_to_drive:\n", " try:\n", @@ -237,12 +168,9 @@ " except Exception as _e:\n", " print(f'(Drive mount failed, using {CACHE_DIR}: {_e})')\n", "\n", - "# A BOUNDED wait. This used to be `while not exists: sleep(5)` with the\n", - "# background job's output discarded, so a failed download was an indefinite\n", - "# hang with nothing on screen -- which is exactly how a HuggingFace 429 looked\n", - "# on 2026-09-17. The job writes .log; a stall raises with the tail of\n", - "# it. A few hundred MB, so 20 minutes before giving up.\n", + "\n", "def _await(sentinel, limit=1200):\n", + " \"\"\"Wait for a background job, with a bound -- a failed download must not hang.\"\"\"\n", " t0 = time.time()\n", " while not os.path.isfile(sentinel):\n", " if time.time() - t0 > limit:\n", @@ -254,10 +182,8 @@ " print(f'{sentinel} \\u2713 ({time.time() - t0:.0f} s)')\n", "\n", "\n", - "_await(STAMP)\n", - "\n", - "if IS_AF3 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", + "if not IS_AF3:\n", + " _await(STAMP)\n", "\n", "print(f'Setup complete! Model: {model}.')\n", "if model == 'chai1':\n", @@ -296,11 +222,7 @@ "#@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. A form field above is a plain assignment, so a notebook\n", - "# run outside Colab -- colab-cli, CI -- has no way to change one. Any field can\n", - "# be set from the environment instead:\n", - "# AF3_NB_OVERRIDES='{\"model\": \"boltz2\", \"msa_mode\": \"single_sequence\"}'\n", - "# Unset, this does nothing at all, which is every interactive run.\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", @@ -318,22 +240,10 @@ "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", - "# THE CCD, for the components this input actually names. It has to happen here\n", - "# and not in the install cell: that cell runs BEFORE you type a ligand, so it\n", - "# can only ever fetch the standard residues -- and a fold naming ATP then dies\n", - "# with `ValueError: Unknown residue type ATP`.\n", - "#\n", - "# NOT build_data, which parses libcifpp's whole components.cif (51,275\n", - "# components, 518 MB) into a 505 MB pickle and costs 48 s of the session.\n", - "# LocalFold's trick instead: pull each component from\n", - "# files.rcsb.org/ligands/download/.cif, kilobytes each. Measured at 0.6 s\n", - "# for 40 components, every field byte-identical to libcifpp's for ALA, SER,\n", - "# GOL, ATP, SEP, NAG, DA and U.\n", - "#\n", - "# Both pickles are written from the SAME fetched set, so they stay\n", - "# self-consistent -- NAG lands in GLYCAN_LINKING_LIGANDS exactly as it would\n", - "# with the full dictionary. A code we did not fetch raises KeyError, which is\n", - "# loud, rather than being silently mis-bonded.\n", + "# The CCD, for the components this input names -- which is why it is here and\n", + "# not in the install cell, which runs before you have typed a ligand. Each one\n", + "# comes from files.rcsb.org (kilobytes, ~0.6 s) instead of build_data parsing\n", + "# libcifpp's whole 518 MB dictionary. A code that was not fetched raises.\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", @@ -447,11 +357,7 @@ "#@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. A form field above is a plain assignment, so a notebook\n", - "# run outside Colab -- colab-cli, CI -- has no way to change one. Any field can\n", - "# be set from the environment instead:\n", - "# AF3_NB_OVERRIDES='{\"model\": \"boltz2\", \"msa_mode\": \"single_sequence\"}'\n", - "# Unset, this does nothing at all, which is every interactive run.\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", @@ -560,10 +466,8 @@ "cmd = ' '.join(cmd)\n", "if run_it:\n", " print(cmd)\n", - " # NOT `!{cmd}`. That reports nothing about how the run ended, and the line\n", - " # below used to print `Done -> ...` whatever happened -- so a hard failure\n", - " # (an ImportError, 230 ms) read as a successful fold with no structures.\n", - " # Popen streams the same output AND yields a status.\n", + " # Popen, not `!{cmd}`: it streams the same output but also yields a status,\n", + " # so a failed run cannot report `Done ->`.\n", " _p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,\n", " stderr=subprocess.STDOUT, text=True, bufsize=1)\n", " for _line in _p.stdout:\n", @@ -681,7 +585,7 @@ "token_chain_ids = conf.get('token_chain_ids', []) # PAE is per-TOKEN\n", "pae = np.array(conf.get('pae', []), dtype=float)\n", "\n", - "# ── Summary (ipTM is None for single-chain jobs — guard before formatting) ─\n", + "# \u2500\u2500 Summary (ipTM is None for single-chain jobs \u2014 guard before formatting) \u2500\n", "def fmt(v):\n", " return f'{v:.3f}' if isinstance(v, (int, float)) else 'n/a'\n", "\n", @@ -693,17 +597,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 — no interface)' if iptm is None else ''))\n", + "print(f'ipTM : {fmt(iptm)}' + (' (single chain \u2014 no interface)' if iptm is None else ''))\n", "print(f'Ranking score : {fmt(summ.get(\"ranking_score\"))}')\n", "print('=' * 38)\n", "\n", - "# ── Plots ───────────────────────────────────────────────────\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", "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 — a line (coloured per chain when there is more than one)\n", + "# pLDDT per residue \u2014 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", @@ -734,7 +638,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 (Å)')\n", + " plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label='PAE (\u00c5)')\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", @@ -776,7 +680,7 @@ "# Instructions \n", "\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", + "**Runtime \u2192 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", @@ -796,12 +700,12 @@ "| `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 — single sequence, no search. |\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** — a ligand or nucleotide raises rather than quietly folding the rest. Templates use the model_1/model_2 parameter sets. |\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) | Requires requesting the weights from Google (non-commercial research only, granted at Google's discretion) — not a direct download. run_alphafold prints the terms at startup. |\n", + "| `alphafold3` | Google DeepMind's own parameters | [AF3 terms of use](https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md) | Requires requesting the weights from Google (non-commercial research only, granted at Google's discretion) \u2014 not a direct download. run_alphafold prints the terms at startup. |\n", "---\n", "\n", "## Input\n", @@ -816,37 +720,37 @@ "| **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… follow AlphaFold 3's canonical\n", - "order (protein → RNA → DNA → ligand). Identical protein sequences are merged, so\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", "`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", + "as typed**. Whitespace and extra colons are forgiven (`SEQ1::::SEQ2` = `SEQ1:SEQ2`) \u2014\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 — RNA/DNA always run\n", + "[ColabFold](https://colabfold.mmseqs.com/) API (protein only \u2014 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–100). |\n", + "| `*.cif` | Best-ranked structure. B-factor = pLDDT (0\u2013100). |\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 × sample combination. |\n", + "| `*_ranking_scores.csv` | Every seed \u00d7 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–90 reliable backbone, 50–70 doubtful, below 50 likely\n", + "pLDDT above 90 is very high, 70\u201390 reliable backbone, 50\u201370 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", + "complex interface, and is `n/a` for a single chain \u2014 there is no interface to score.\n", "\n", "## Troubleshooting\n", "\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", + "**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", "`single_sequence`. **Download popup blocked**: disable your ad blocker.\n", "\n", "## Licence\n", From 1b4fe8d12b7f611c318d4a61e1ad0e4a5603b4d8 Mon Sep 17 00:00:00 2001 From: Sergey Ovchinnikov Date: Thu, 17 Sep 2026 03:45:59 +0000 Subject: [PATCH 07/11] trim the rest of the commentary in the three code cells Still too much explaining of history rather than of code: what %%time did, what a 429 looked like, what `Done ->` used to print, the exact RMSDs a language model is worth. What is left says what a line does or why it could not be simpler. No behaviour change; the AlphaFold 3 message still names the request page and the directory to put the file in. --- ColabFold2_preview.ipynb | 75 +++++++++++++--------------------------- 1 file changed, 24 insertions(+), 51 deletions(-) diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb index 4dfb193bc..a66831141 100644 --- a/ColabFold2_preview.ipynb +++ b/ColabFold2_preview.ipynb @@ -45,7 +45,6 @@ "outputs": [], "source": [ "#@title Install dependencies (~35 s)\n", - "# No `%%time` here: a cell magic must be the first line and `#@title` already is.\n", "import os, time, glob, shutil, sys\n", "_T0 = time.time()\n", "\n", @@ -56,8 +55,8 @@ "#@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", - "# Form fields are plain assignments, so a headless run (colab-cli, CI) sets them\n", - "# through the environment: AF3_NB_OVERRIDES='{\"model\": \"boltz2\"}'\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", @@ -69,29 +68,27 @@ "AF2_DIR = 'af2_params'\n", "IS_AF3 = (model == 'alphafold3')\n", "IS_AF2 = model.startswith('af2_')\n", - "# int8: the same weights stored 8-bit and expanded on load, which is what keeps\n", - "# the download to a few hundred MB. AF2 and AF3 ship their own float32 files.\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", "\n", "def _sh(cmd, what):\n", - " \"\"\"Run a shell command and raise if it fails -- os.system's status is easy to drop.\"\"\"\n", + " \"\"\"Run a shell command, raising if it fails.\"\"\"\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", " print('Installing packages...')\n", - " # --no-deps throughout: Colab already ships jax and the CUDA stack, and letting\n", - " # pip re-resolve them re-downloads gigabytes. So every third-party import the\n", - " # package needs is listed here instead.\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", " \"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", " '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: a wheel published minutes ago may not be in PyPI's index yet.\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", " break\n", @@ -104,13 +101,12 @@ " 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 # the only real proof the install worked\n", + " import alphafold3 # confirms the install before anything depends on it\n", " os.system('touch ALPHAFOLD3_READY')\n", " print(f'Packages installed ({alphafold3.__file__}).')\n", "\n", - "# tokamax enables its Triton kernels for every GPU with cc >= 8.0, but they need\n", - "# more shared memory than Ada cards have and fail at launch. Restrict them to\n", - "# datacenter GPUs (A100 cc 8.0, H100 cc 9.0+); Ada/L4 use XLA like a T4 does.\n", + "# tokamax's Triton kernels need more shared memory than Ada cards have, so\n", + "# restrict them to datacenter GPUs (A100 cc 8.0, H100 cc 9.0+).\n", "try:\n", " import tokamax\n", " _gu = os.path.join(os.path.dirname(tokamax.__file__), '_src', 'gpu_utils.py')\n", @@ -124,15 +120,10 @@ "except Exception as _e:\n", " print(f'(tokamax patch skipped: {_e})')\n", "\n", - "# Weights, fetched in the background by the same code the run uses, so the cache\n", - "# layout cannot drift between the two.\n", + "# Weights, fetched in the background by the same code the run uses.\n", "STAMP = f'WEIGHTS_DONE_{model}_{PRECISION}'\n", "if IS_AF3:\n", - " # NOT downloaded. DeepMind's AlphaFold 3 parameters are granted on request,\n", - " # for non-commercial research, at Google's discretion -- they are not ours to\n", - " # fetch on your behalf. Apply at\n", - " # https://docs.google.com/forms/d/e/1FAIpQLSfWZAgo1aYk0O4MuAXZj8xRQ8DafeFJnldNOnh_13qAx2ceZw/viewform\n", - " # and put the file you are given in af3_native_weights/.\n", + " # DeepMind grants these on request; supply your own copy.\n", " os.makedirs(NATIVE_DIR, exist_ok=True)\n", " _blobs = glob.glob(f'{NATIVE_DIR}/*.bin.zst')\n", " if not _blobs:\n", @@ -155,8 +146,7 @@ " '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 on a small\n", - "# input); Drive survives. Falls back to /tmp rather than failing.\n", + "# /tmp is wiped with the VM, so a fresh session recompiles (~53 s); Drive survives.\n", "CACHE_DIR = '/tmp/af3_cache'\n", "if persist_cache_to_drive:\n", " try:\n", @@ -170,7 +160,7 @@ "\n", "\n", "def _await(sentinel, limit=1200):\n", - " \"\"\"Wait for a background job, with a bound -- a failed download must not hang.\"\"\"\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", @@ -240,10 +230,8 @@ "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", - "# The CCD, for the components this input names -- which is why it is here and\n", - "# not in the install cell, which runs before you have typed a ligand. Each one\n", - "# comes from files.rcsb.org (kilobytes, ~0.6 s) instead of build_data parsing\n", - "# libcifpp's whole 518 MB dictionary. A code that was not fetched raises.\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", @@ -347,7 +335,6 @@ "outputs": [], "source": [ "#@title Run the model\n", - "# No `%%time` -- see the install cell.\n", "import os, shutil, subprocess, glob, time\n", "_T0 = time.time()\n", "\n", @@ -377,10 +364,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", @@ -402,21 +386,17 @@ " nojit = True\n", " print('No GPU detected - running on CPU with XLA attention + --nojit (slow, but avoids the compile).')\n", "elif cap < 8.0:\n", - " # T4 / V100 (cc 7.x): XLA attention; disable the custom-kernel fusion pass.\n", - " # (Triton GEMM is not supported on these cards, so it is not disabled here.)\n", + " # T4 / V100: XLA attention, and no custom-kernel fusion pass.\n", " flash_impl = 'xla'\n", " xla_flags = ['--xla_disable_hlo_passes=custom-kernel-fusion-rewriter']\n", " print(f'Pre-Ampere GPU (compute capability {cap}) - XLA attention + custom-kernel fusion disabled.')\n", "elif 8.0 < cap < 9.0:\n", - " # L4 / Ada / consumer Ampere (cc 8.6 / 8.9): limited shared memory. XLA's Triton GEMM\n", - " # kernels exceed it ('Shared memory size limit exceeded'), so disable Triton GEMM\n", - " # (falls back to cuBLAS) and use XLA attention to also avoid the Triton attention kernel.\n", + " # L4 / Ada: limited shared memory, so no Triton kernels -- XLA and cuBLAS.\n", " flash_impl = 'xla'\n", " xla_flags = ['--xla_gpu_enable_triton_gemm=false']\n", " print(f'Ada/consumer GPU (compute capability {cap}) - XLA attention + Triton GEMM disabled (shared-memory limit).')\n", "else:\n", - " # A100 (cc 8.0) and H100 (cc 9.0+): ample shared memory. Triton flash attention,\n", - " # with Triton GEMM disabled per AlphaFold 3's recommended XLA_FLAGS.\n", + " # A100 / H100: Triton flash attention, Triton GEMM off per AlphaFold 3.\n", " flash_impl = 'triton'\n", " xla_flags = ['--xla_gpu_enable_triton_gemm=false']\n", " print(f'Datacenter GPU (compute capability {cap}) - Triton flash attention + Triton GEMM disabled.')\n", @@ -431,10 +411,7 @@ "\n", "print('XLA_FLAGS =', os.environ.get('XLA_FLAGS', '(unset)'))\n", "\n", - "# Weights. Every ported model resolves its own cache directory (populated by the\n", - "# install cell), so --model_dir is passed only for the two whose parameters come\n", - "# from DeepMind directly: AlphaFold 3's, and AlphaFold 2's (CC BY 4.0, fetched\n", - "# into af2_params by the install cell).\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", @@ -451,11 +428,8 @@ "]\n", "if msa_mode == 'mmseqs2_server':\n", " cmd.append('--use_msa_server')\n", - "# chai-1 folds from ESM2 and ESMFold2 from ESM-C; without it they are a\n", - "# different model, not a slightly worse one (a natural protein goes to 5.70 A\n", - "# where chai-1 reaches 0.642, and an ESMFold2 variant with no MSA encoder has\n", - "# nothing left to fold from at all). Both towers run in-process and download on\n", - "# demand, which is why run_alphafold makes it opt-in and this passes it.\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", @@ -466,8 +440,7 @@ "cmd = ' '.join(cmd)\n", "if run_it:\n", " print(cmd)\n", - " # Popen, not `!{cmd}`: it streams the same output but also yields a status,\n", - " # so a failed run cannot report `Done ->`.\n", + " # Popen rather than `!`: streams the output and gives an exit status.\n", " _p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,\n", " stderr=subprocess.STDOUT, text=True, bufsize=1)\n", " for _line in _p.stdout:\n", From 116fe28f0cfec5142b3a859bdfd1d097f70e3389 Mon Sep 17 00:00:00 2001 From: Sergey Ovchinnikov Date: Thu, 17 Sep 2026 03:55:11 +0000 Subject: [PATCH 08/11] alphafold3's weights are a public release, so fetch them again The table row said they are granted on request and the install cell had been changed to match; the bucket is DeepMind's intended public release, so the row is what needed fixing. A copy already in af3_native_weights/ is now used as-is rather than deleted before downloading. Verified on a cold Colab T4: 1 GB in 15 s, fold in 62 s. --- ColabFold2_preview.ipynb | 34 ++++++++++++++++++---------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb index a66831141..2155fdd59 100644 --- a/ColabFold2_preview.ipynb +++ b/ColabFold2_preview.ipynb @@ -65,6 +65,7 @@ "\n", "VERSION = '3.1.10' # 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", @@ -122,20 +123,19 @@ "\n", "# Weights, fetched in the background by the same code the run uses.\n", "STAMP = f'WEIGHTS_DONE_{model}_{PRECISION}'\n", - "if IS_AF3:\n", - " # DeepMind grants these on request; supply your own copy.\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", - " _blobs = glob.glob(f'{NATIVE_DIR}/*.bin.zst')\n", - " if not _blobs:\n", - " raise RuntimeError(\n", - " f'No AlphaFold 3 parameters found in {NATIVE_DIR}/.\\n'\n", - " 'DeepMind grants these on request (non-commercial research); this '\n", - " 'notebook cannot download them for you. Request them at\\n'\n", - " ' https://github.com/google-deepmind/alphafold3#obtaining-model-parameters\\n'\n", - " f'then upload the .bin.zst file into {NATIVE_DIR}/ and re-run this cell.\\n'\n", - " 'Every other model in the dropdown downloads its own weights.')\n", - " print(f'Using your AlphaFold 3 parameters: {_blobs[0]}')\n", - "elif not os.path.isfile(STAMP):\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", @@ -172,8 +172,10 @@ " print(f'{sentinel} \\u2713 ({time.time() - t0:.0f} s)')\n", "\n", "\n", - "if not IS_AF3:\n", - " _await(STAMP)\n", + "_await(STAMP)\n", + "\n", + "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", @@ -678,7 +680,7 @@ "| `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) | Requires requesting the weights from Google (non-commercial research only, granted at Google's discretion) \u2014 not a direct download. run_alphafold prints the terms at startup. |\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", "---\n", "\n", "## Input\n", From a9142020e2358eebb10998f96b513a8c20cf480d Mon Sep 17 00:00:00 2001 From: Sergey Ovchinnikov Date: Fri, 18 Sep 2026 11:34:28 +0000 Subject: [PATCH 09/11] ColabFold2_preview: alphafold3-colabfold 3.1.11 The int8 weights were republished with per-block scales (and with the parameters several blobs were missing -- boltz2 and opendde would not load at all). 3.1.10's loader raises on the new scale records, so the pin has to move with them. Co-Authored-By: Claude Opus 5 --- ColabFold2_preview.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ColabFold2_preview.ipynb b/ColabFold2_preview.ipynb index 2155fdd59..462a5b4b0 100644 --- a/ColabFold2_preview.ipynb +++ b/ColabFold2_preview.ipynb @@ -63,7 +63,7 @@ " globals()[_k] = _v\n", " print(f'override: {_k} = {_v!r}')\n", "\n", - "VERSION = '3.1.10' # package and run_alphafold.py both come from this tag\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", From a56df9ca4501841679275ed05af3971a337b6180 Mon Sep 17 00:00:00 2001 From: Ronit Baldaniya <222511871+RonitBStudent@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:07:52 -0500 Subject: [PATCH 10/11] Add --version flag to colabfold_batch, colabfold_search and colabfold_split_msas Reproducible workflows need to record the tool version, but none of the CLIs exposed one. Add a standard argparse `--version` action to all three entry points, backed by a shared `get_version()` helper that reports the installed distribution version plus the git commit when installed from VCS (the same string colabfold_batch already logs at startup). Closes #414 --- colabfold/batch.py | 7 +++---- colabfold/mmseqs/search.py | 3 ++- colabfold/mmseqs/split_msas.py | 3 +++ colabfold/utils.py | 9 +++++++++ 4 files changed, 17 insertions(+), 5 deletions(-) diff --git a/colabfold/batch.py b/colabfold/batch.py index e0359d073..3b41be63d 100644 --- a/colabfold/batch.py +++ b/colabfold/batch.py @@ -68,6 +68,7 @@ NO_GPU_FOUND, CIF_REVISION_DATE, get_commit, + get_version, setup_logging, CFMMCIFIO, AF3Utils, @@ -1814,6 +1815,7 @@ def generate_af3_input( def main(): parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter) + parser.add_argument("--version", action="version", version=f"%(prog)s {get_version()}") parser.add_argument( "input", default="input", @@ -2233,10 +2235,7 @@ def comma_separated_list(arg_string): setup_logging(Path(args.results).joinpath("log.txt"), verbose=args.debug_logging) - version = importlib_metadata.version("colabfold") - commit = get_commit() - if commit: - version += f" ({commit})" + version = get_version() logger.info(f"Running colabfold {version}") diff --git a/colabfold/mmseqs/search.py b/colabfold/mmseqs/search.py index 9c3be7f03..530be8902 100644 --- a/colabfold/mmseqs/search.py +++ b/colabfold/mmseqs/search.py @@ -14,7 +14,7 @@ from typing import List, Union from colabfold.input import get_queries, msa_to_str, safe_filename -from colabfold.utils import AF3Utils +from colabfold.utils import AF3Utils, get_version logger = logging.getLogger(__name__) @@ -291,6 +291,7 @@ def mmseqs_search_pair( def main(): parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter) + parser.add_argument("--version", action="version", version=f"%(prog)s {get_version()}") parser.add_argument( "query", type=Path, diff --git a/colabfold/mmseqs/split_msas.py b/colabfold/mmseqs/split_msas.py index b0e6e67f8..5c0eb80da 100644 --- a/colabfold/mmseqs/split_msas.py +++ b/colabfold/mmseqs/split_msas.py @@ -8,6 +8,8 @@ from tqdm import tqdm +from colabfold.utils import get_version + logger = logging.getLogger(__name__) @@ -36,6 +38,7 @@ def main(): parser = ArgumentParser( description="Take an a3m database from the colabdb search and turn it into a folder of a3m files" ) + parser.add_argument("--version", action="version", version=f"%(prog)s {get_version()}") parser.add_argument( "search_folder", help="The search folder in which you ran colabfold_search with the final.a3m", diff --git a/colabfold/utils.py b/colabfold/utils.py index 0d3314f21..1345fa216 100644 --- a/colabfold/utils.py +++ b/colabfold/utils.py @@ -77,6 +77,15 @@ def get_commit() -> Optional[str]: return direct_url["vcs_info"]["commit_id"] +def get_version() -> str: + """Installed colabfold version, with the git commit appended when installed from VCS.""" + version = distribution("colabfold").version + commit = get_commit() + if commit: + version += f" ({commit})" + return version + + # Copied from Bio.PDB to override _save_dict method # https://github.com/biopython/biopython/blob/biopython-179/Bio/PDB/mmcifio.py # We add poly_seq and revision_date so that AF2 can read these cif files From f95293b8d71a1688b0fd2908d5d0009fc87de189 Mon Sep 17 00:00:00 2001 From: Ronit Baldaniya <222511871+RonitBStudent@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:43:30 -0500 Subject: [PATCH 11/11] Make colabfold_batch --version work without the alphafold extra colabfold.batch raises at import time when the `alphafold` extra is not installed, so argparse never got to see `--version` in a base install. Answer the flag in that guard before failing, so the version is always reachable. --- colabfold/batch.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/colabfold/batch.py b/colabfold/batch.py index 3b41be63d..7d0d0bfcf 100644 --- a/colabfold/batch.py +++ b/colabfold/batch.py @@ -38,6 +38,14 @@ try: import alphafold except ModuleNotFoundError: + if "--version" in sys.argv[1:]: + # `colabfold_batch --version` must still answer in a base install + # (without the `alphafold` extra). argparse would handle the flag, but + # main() is never reached because the imports below fail first. + from colabfold.utils import get_version + + print(f"{os.path.basename(sys.argv[0])} {get_version()}") + sys.exit(0) raise RuntimeError( "\n\nalphafold is not installed. Please run `pip install colabfold[alphafold]`\n" )