Integration of RGI-toolkit

Boltz1.ipynb

@@ -70,6 +70,11 @@
         "#@markdown - Use `:` to specify multiple DNA sequences\n",
         "jobname = 'test'  #@param {type:\"string\"}\n",
         "\n",
+        "msa_mode = \"mmseqs2_server\" #@param [\"mmseqs2_server\", \"single_sequence\"]\n",
+        "import json\n",
+        "for _k, _v in json.loads(os.environ.get('BOLTZ_NB_OVERRIDES', '{}')).items():\n",
+        "    if _k in globals(): globals()[_k] = _v\n",
+        "\n",
         "# Clean up the query sequence and jobname\n",
         "query_sequence = \"\".join(query_sequence.split())\n",
         "ligand_input = \"\".join(ligand_input.split())\n",
@@ -219,7 +224,11 @@
         "# Optionally, print the output for verification\n",
         "#print(f\"Generated FASTA file '{queries_fasta}':\\n\")\n",
         "#for header, sequence in fasta_entries:\n",
-        "#    print(f\"{header}\\n{sequence}\\n\")\n"
+        "#    print(f\"{header}\\n{sequence}\\n\")\n",
+        "\n",
+        "print('Chain IDs for restraints:')\n",
+        "for header, sequence in fasta_entries:\n",
+        "    print(header.split('|')[0].lstrip('>'), header.split('|')[1], len(sequence) if header.split('|')[1] == 'protein' else '')\n"
       ],
       "metadata": {
         "cellView": "form",
@@ -228,6 +237,60 @@
       "execution_count": null,
       "outputs": []
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rgi-restraints",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title RGI: Making restraint-guided protein folding inference accessible to all! (optional)\n",
+        "#@markdown Fill in this form before **Runtime → Run all**. Leave `use_rgi` off for vanilla. Leave unused restraint\n",
+        "#@markdown fields empty.\n",
+        "use_rgi = False #@param {type:\"boolean\"}\n",
+        "#@markdown **distance** — enter two atom selections and their target distance in Å. Example: `chain A and resid\n",
+        "#@markdown 1 to 10` and `chain A and resid 40 to 50`, with `target_distance = 25`.\n",
+        "distance_atom_selection1 = \"\" #@param {type:\"string\"}\n",
+        "distance_atom_selection2 = \"\" #@param {type:\"string\"}\n",
+        "target_distance = 25 #@param {type:\"raw\"}\n",
+        "#@markdown **conformer** — enter chain IDs such as `B,C`, or `ligands` for all ligand chains. Leave empty to\n",
+        "#@markdown disable.\n",
+        "conformer_chains = \"\" #@param {type:\"string\"}\n",
+        "#@markdown **angle** — enter three atom selections and the target in degrees. Selection 2 is the vertex. Leave\n",
+        "#@markdown all three empty to disable.\n",
+        "angle_atom_selection1 = \"\" #@param {type:\"string\"}\n",
+        "angle_atom_selection2 = \"\" #@param {type:\"string\"}\n",
+        "angle_atom_selection3 = \"\" #@param {type:\"string\"}\n",
+        "target_angle = 90 #@param {type:\"raw\"}\n",
+        "#@markdown **custom** — set named selections and an energy expression, e.g. `harmonic(distance(A, B), 25.0)`.\n",
+        "#@markdown Leave the energy empty to disable.\n",
+        "custom_selections = {\"A\": \"chain A and resid 1 to 10\", \"B\": \"chain A and resid 40 to 50\"} #@param {type:\"raw\"}\n",
+        "custom_energy = \"\" #@param {type:\"string\"}\n",
+        "#@markdown **RMSD** — upload a reference using Colab's **Files → Upload** and enter its filename in either\n",
+        "#@markdown `ref_pdb` or `ref_cif`. Leave both filenames empty to disable. Empty target/ref selections mean the\n",
+        "#@markdown whole structure; e.g. `chain A and name CA` selects C-alpha atoms.\n",
+        "ref_pdb = \"\" #@param {type:\"string\"}\n",
+        "ref_cif = \"\" #@param {type:\"string\"}\n",
+        "atom_selection_target = \"\" #@param {type:\"string\"}\n",
+        "atom_selection_ref = \"\" #@param {type:\"string\"}\n",
+        "target_rmsd = 0 #@param {type:\"raw\"}\n",
+        "#@markdown **Multiple restraints:** enter lists in the same fields, e.g. `distance_atom_selection2 = [\"chain A\n",
+        "#@markdown and resid 20\", \"chain A and resid 40\"]` and `target_distance = [15, 25]`. A single selection/value is\n",
+        "#@markdown shared across the entries. All non-single lists must have the same length.\n",
+        "#@markdown **Selections:** all fields accept the full RGI-toolkit selection language. Residues start at 1 within\n",
+        "#@markdown each chain. The input cell prints the chain IDs. Distance and angle use atom-group centroids.\n",
+        "#@markdown **Advanced (optional)** — an ordinary RGI-toolkit dict for additional restraints, flat-bottomed\n",
+        "#@markdown penalties, weights, activation windows or conformer settings. Lists append to the fields above. Use\n",
+        "#@markdown `{\"config_path\": \"restraints.yaml\"}` to read a file. [Examples and native config\n",
+        "#@markdown reference](https://github.com/cddlab/rgi_toolkit/blob/main/docs/colabfold.md).\n",
+        "restraints_config = {} #@param {type:\"raw\"}\n",
+        "#@markdown After changing this form, rerun this cell and the prediction cell, or use **Run all**. Settings are\n",
+        "#@markdown saved in the notebook itself.\n",
+        "#@markdown When switching `use_rgi`, rerun **Install dependencies** too to select the matching Boltz environment.\n"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -238,34 +301,39 @@
       "outputs": [],
       "source": [
         "#@title Install dependencies\n",
-        "%%time\n",
-        "import os\n",
-        "if not os.path.isfile(\"COLABFOLD_READY\"):\n",
-        "  print(\"installing colabfold...\")\n",
-        "  os.system(\"pip install -q --no-warn-conflicts 'colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold'\")\n",
-        "  if os.environ.get('TPU_NAME', False) != False:\n",
-        "    os.system(\"pip uninstall -y jax jaxlib\")\n",
-        "    os.system(\"pip install --no-warn-conflicts --upgrade dm-haiku==0.0.10 'jax[cuda12_pip]'==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\")\n",
-        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabfold colabfold\")\n",
-        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/alphafold alphafold\")\n",
-        "  # hack to fix TF crash\n",
-        "  os.system(\"rm -f /usr/local/lib/python3.*/dist-packages/tensorflow/core/kernels/libtfkernel_sobol_op.so /usr/local/lib/python3.*/dist-packages/tensorflow/lite/python/*/*.so\")\n",
-        "  os.system(\"touch COLABFOLD_READY\")\n",
-        "\n",
-        "if not os.path.isfile(\"BOLZ_READY\"):\n",
-        "  os.system(\"apt-get install -y aria2\")\n",
-        "  os.system(\"pip install -q --no-warn-conflicts boltz\")\n",
-        "  os.system(\"mkdir weights\")\n",
-        "  os.system(\"aria2c -d weights -x8 -s8 https://colabfold.steineggerlab.workers.dev/boltz1.ckpt\")\n",
-        "  os.system(\"aria2c -d weights -x8 -s8 https://colabfold.steineggerlab.workers.dev/ccd.pkl\")\n",
-        "  os.system(\"touch BOLZ_READY\")"
+        "import shutil, subprocess, sys\n",
+        "from pathlib import Path\n",
+        "use_rgi = globals().get('use_rgi', False)\n",
+        "if shutil.which('uv') is None:\n",
+        "    subprocess.run([sys.executable, '-m', 'pip', 'install', '-q', 'uv'], check=True)\n",
+        "_colabfold_source = os.environ.get('RGI_COLABFOLD_SOURCE', 'git+https://github.com/sokrypton/ColabFold')\n",
+        "_toolkit_source = os.environ.get('RGI_TOOLKIT_SOURCE', 'git+https://github.com/cddlab/rgi_toolkit.git@054cd4c81a7f198df12fdd72f88d5caa7f8d9615')\n",
+        "subprocess.run(['uv', 'pip', 'install', '--python', sys.executable, _colabfold_source], check=True)\n",
+        "# Boltz requires Python 3.12; isolate it from Colab's Python and NumPy versions.\n",
+        "boltz_env = Path('.cache') / ('boltz-rgi' if use_rgi else 'boltz-vanilla')\n",
+        "if not (boltz_env / 'bin/python').exists():\n",
+        "    subprocess.run(['uv', 'venv', '--python', '3.12', str(boltz_env)], check=True)\n",
+        "boltz_python = str(boltz_env / 'bin/python')\n",
+        "_boltz_source = ('git+https://github.com/cddlab/boltz_restr.git@rgi-integration' if use_rgi else 'boltz==2.2.1')\n",
+        "subprocess.run(['uv', 'pip', 'install', '--python', boltz_python, _boltz_source], check=True)\n",
+        "if use_rgi:\n",
+        "    subprocess.run(['uv', 'pip', 'install', '--python', sys.executable, _toolkit_source], check=True)\n",
+        "    subprocess.run(['uv', 'pip', 'install', '--python', boltz_python, _toolkit_source], check=True)\n"
       ]
     },
     {
       "cell_type": "code",
       "source": [
         "#@title Generate MSA with ColabFold\n",
-        "!colabfold_batch \"{queries_path}\" \"{jobname}\" --msa-only"
+        "if msa_mode == 'mmseqs2_server' and csv_entries:\n",
+        "    from colabfold.colabfold import run_mmseqs2\n",
+        "    for seq_id, seq in csv_entries:\n",
+        "        msa_path = Path(jobname) / f'{seq_id}.a3m'\n",
+        "        if not msa_path.exists():\n",
+        "            a3m = run_mmseqs2([seq], str(Path(jobname) / (seq_id + '_msa')), use_env=True)[0]\n",
+        "            msa_path.write_text(a3m)\n",
+        "else:\n",
+        "    print('Single sequence: skipping the MSA server.')\n"
       ],
       "metadata": {
         "cellView": "form",
@@ -278,7 +346,46 @@
       "cell_type": "code",
       "source": [
         "#@title Predict structure using boltz\n",
-        "!boltz predict --out_dir \"{jobname}\" \"{jobname}/{jobname}.fasta\" --cache weights"
+        "import yaml\n",
+        "num_recycles = 3 #@param {type:\"integer\"}\n",
+        "num_diffusion_samples = 1 #@param {type:\"integer\"}\n",
+        "seed = 42 #@param {type:\"integer\"}\n",
+        "# Read the standard Colab fields before starting inference.\n",
+        "if globals().get('use_rgi', False):\n",
+        "    from rgi_toolkit.notebook_colab import config_from_fields\n",
+        "    rgi_config, rgi_conformer_chains = config_from_fields(globals())\n",
+        "else:\n",
+        "    use_rgi = False\n",
+        "    rgi_config, rgi_conformer_chains = None, ''\n",
+        "\n",
+        "from colabfold.rgi.boltz import notebook_input\n",
+        "\n",
+        "# Use YAML so the predictor can read the shared RGI config and entity opt-ins.\n",
+        "boltz_input = notebook_input(fasta_entries, config=rgi_config if use_rgi else None,\n",
+        "    conformer_chains=rgi_conformer_chains, single_sequence=(msa_mode == 'single_sequence'))\n",
+        "input_yaml = Path(jobname) / f'{jobname}.yaml'\n",
+        "input_yaml.write_text(yaml.safe_dump(boltz_input, sort_keys=False))\n",
+        "# Refresh RGI features without changing vanilla-only cache reuse.\n",
+        "boltz_output = Path(jobname) / f'boltz_results_{input_yaml.stem}'\n",
+        "rgi_cache_marker = boltz_output / '.rgi-enabled'\n",
+        "if use_rgi or rgi_cache_marker.exists():\n",
+        "    import shutil\n",
+        "    if boltz_output.exists():\n",
+        "        shutil.rmtree(boltz_output)\n",
+        "    if use_rgi:\n",
+        "        boltz_output.mkdir(parents=True, exist_ok=True)\n",
+        "        rgi_cache_marker.touch()\n",
+        "subprocess.run([boltz_python, '-m', 'boltz.main', 'predict', str(input_yaml),\n",
+        "    '--model', 'boltz1', '--out_dir', jobname, '--cache', str(Path('.cache') / 'boltz-weights'),\n",
+        "    '--recycling_steps', str(num_recycles), '--diffusion_samples', str(num_diffusion_samples),\n",
+        "    '--seed', str(seed), '--override'], check=True)\n",
+        "import glob\n",
+        "if not glob.glob(f'{jobname}/boltz_results_{jobname}/predictions/{jobname}/*.cif'):\n",
+        "    raise RuntimeError('Prediction produced no structure files. Check the log above.')\n",
+        "if use_rgi:\n",
+        "    from colabfold.citations import write_rgi_citation\n",
+        "    write_rgi_citation(boltz_output)\n",
+        "print('RGI ON: check nonzero built-spec counts and final energies above.' if use_rgi else 'Vanilla prediction complete.')\n"
       ],
       "metadata": {
         "cellView": "form",
@@ -307,7 +414,8 @@
         "with zipfile.ZipFile(zip_filename, 'w') as zipf:\n",
         "    coverage_png_files = glob.glob(os.path.join(jobname, '*_coverage.png'))\n",
         "    a3m_files = glob.glob(os.path.join(jobname, '*.a3m'))\n",
-        "    for file in coverage_png_files + a3m_files:\n",
+        "    bibtex_files = glob.glob(os.path.join(jobname, f'boltz_results_{jobname}', 'cite.bibtex'))\n",
+        "    for file in coverage_png_files + a3m_files + [str(input_yaml)] + bibtex_files:\n",
         "        arcname = os.path.basename(file)  # Use only the file name\n",
         "        zipf.write(file, arcname=arcname)\n",
         "\n",
@@ -334,4 +442,4 @@
       "outputs": []
     }
   ]
-}
+}