Update relax notebook to expose more parameters

sokrypton · Nov 9, 2023 · 94ddcfb · 94ddcfb
1 parent 29395e7
commit 94ddcfb
Showing 1 changed file with 42 additions and 14 deletions.
diff --git a/beta/relax_amber.ipynb b/beta/relax_amber.ipynb
@@ -28,19 +28,29 @@
     {
       "cell_type": "markdown",
       "source": [
-        "#relax your structure (using amber)"
+        "\n",
+        "<img src=\"https://raw.githubusercontent.com/sokrypton/ColabFold/main/.github/ColabFold_Marv_Logo_Small.png\" height=\"200\" align=\"right\" style=\"height:240px\">\n",
+        "\n",
+        "# Relax your structure using OpenMM/Amber\n",
+        "\n",
+        "Easy to use protein structure relaxation using [AlphaFold2](https://www.nature.com/articles/s41586-021-03819-2) and [OpenMM](https://github.com/openmm/openmm).\n",
+        "\n",
+        "This notebooks is part of ColabFold. See [GitHub](https://github.com/sokrypton/ColabFold) for other notebooks.\n",
+        "\n",
+        "[Mirdita M, Schütze K, Moriwaki Y, Heo L, Ovchinnikov S, Steinegger M. ColabFold: Making protein folding accessible to all.\n",
+        "*Nature Methods*, 2022](https://www.nature.com/articles/s41592-022-01488-1)\n"
       ],
       "metadata": {
-        "id": "TXSecRRnpGeN"
+        "id": "rvfNfOLaE4bW"
       }
     },
     {
       "cell_type": "code",
       "source": [
         "%%time\n",
-        "#@title setup\n",
+        "#@title Setup AlphaFold2 and OpenMM\n",
         "import sys, os\n",
-        "from sys import version_info \n",
+        "from sys import version_info\n",
         "PYTHON_VERSION = f\"{version_info.major}.{version_info.minor}\"\n",
         "\n",
         "if not os.path.isfile(\"ALPHAFOLD_READY\"):\n",
@@ -51,8 +61,8 @@
         "\n",
         "if not os.path.isfile(\"CONDA_READY\"):\n",
         "  print(\"installing conda...\")\n",
-        "  os.system(\"wget -qnc https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh\")\n",
-        "  os.system(\"bash Mambaforge-colab-Linux-x86_64.sh -bfp /usr/local\")\n",
+        "  os.system(\"wget -qnc https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh\")\n",
+        "  os.system(\"bash Mambaforge-Linux-x86_64.sh -bfp /usr/local\")\n",
         "  os.system(\"mamba config --set auto_update_conda false\")\n",
         "  os.system(\"touch CONDA_READY\")\n",
         "\n",
@@ -131,16 +141,17 @@
         "        lines.append(line)\n",
         "  return \"\\n\".join(lines)\n",
         "\n",
-        "def relax_me(pdb_in, pdb_out):\n",
+        "def relax_me(pdb_in, pdb_out, max_iterations, tolerance, stiffness, use_gpu):\n",
         "  pdb_str = pdb_to_string(pdb_in)\n",
         "  protein_obj = protein.from_pdb_string(pdb_str)\n",
         "  amber_relaxer = relax.AmberRelaxation(\n",
-        "    max_iterations=0,\n",
-        "    tolerance=2.39,\n",
-        "    stiffness=10.0,\n",
+        "    max_iterations=max_iterations,\n",
+        "    tolerance=tolerance,\n",
+        "    stiffness=stiffness,\n",
         "    exclude_residues=[],\n",
         "    max_outer_iterations=3,\n",
-        "    use_gpu=False)\n",
+        "    use_gpu=use_gpu\n",
+        "  )\n",
         "  relaxed_pdb_lines, _, _ = amber_relaxer.process(prot=protein_obj)\n",
         "  with open(pdb_out, 'w') as f:\n",
         "      f.write(relaxed_pdb_lines)"
@@ -155,15 +166,32 @@
     {
       "cell_type": "code",
       "source": [
-        "#@title #RELAX\n",
+        "#@title #Relax\n",
         "#@markdown - Click the little ▶ play icon to the left to get an upload prompt.\n",
         "#@markdown - After relax is done, `relaxed.pdb` will automatically download.\n",
         "#@markdown - If download was blocked, click the little folder 📁 icon on the left, right-click `relaxed.pdb` and select Download!\n",
+        "#@markdown ----\n",
+        "#@markdown ## Parameters\n",
+        "max_iterations = 2000 #@param {type:\"slider\", min:0, max:5000, step:250}\n",
+        "#@markdown - Low values might never converge and crash and high values might also run for a long time. `0` = infinity is the AF2 default, however it might run for a long time.\n",
+        "tolerance = 2.39 # @param {type:\"number\"}\n",
+        "#@markdown - kcal/mol, the energy tolerance of L-BFGS.\n",
+        "stiffness = 10.0 # @param {type:\"number\"}\n",
+        "#@markdown - kcal/mol A^2, spring constant of heavy atom restraining potential.\n",
+        "#@markdown - **Note:** Descriptions and default values are taken AF2.\n",
+        "use_gpu = False # @param {type:\"boolean\"}\n",
+        "#@markdown - set \"Runtime->Change runtime type\" to a GPU in Google Colab if you activate this.\n",
         "\n",
         "from google.colab import files\n",
         "pdb_dict = files.upload()\n",
-        "relax_me(pdb_in=list(pdb_dict.keys())[0],\n",
-        "         pdb_out=\"relaxed.pdb\")\n",
+        "relax_me(\n",
+        "    pdb_in=list(pdb_dict.keys())[0],\n",
+        "    pdb_out=\"relaxed.pdb\",\n",
+        "    max_iterations=max_iterations,\n",
+        "    tolerance=tolerance,\n",
+        "    stiffness=stiffness,\n",
+        "    use_gpu=use_gpu\n",
+        ")\n",
         "files.download(f'relaxed.pdb')"
       ],
       "metadata": {