Merge pull request #285 from vaclavhanzl/fix-colab-notebook

Fix notebook: Colab now has python 3.8, fix imports, mitigate UTF-8 glitch

notebooks/OpenFold.ipynb

@@ -121,10 +121,11 @@
         "    %env PATH=/opt/conda/bin:{PATH}\n",
         "\n",
         "    # Install the required versions of all dependencies.\n",
+        "    %shell conda install -y -q conda==4.13.0\n",
         "    %shell conda install -y -q -c conda-forge -c bioconda \\\n",
         "      kalign2=2.04 \\\n",
         "      hhsuite=3.3.0 \\\n",
-        "      python=3.7 \\\n",
+        "      python=3.8 \\\n",
         "      2>&1 1>/dev/null\n",
         "    %shell pip install -q \\\n",
         "      ml-collections==0.1.0 \\\n",
@@ -180,15 +181,12 @@
         "    %shell cp -f /content/stereo_chemical_props.txt /content/openfold/openfold/resources\n",
         "    %shell /usr/bin/python3 -m pip install -q ./openfold\n",
         "\n",
-        "    if(relax_prediction):\n",
-        "      %shell conda install -y -q -c conda-forge \\\n",
-        "        openmm=7.5.1 \\\n",
-        "        pdbfixer=1.7\n",
-        "      \n",
-        "      # Apply OpenMM patch.\n",
-        "      %shell pushd /opt/conda/lib/python3.7/site-packages/ && \\\n",
-        "          patch -p0 < /content/openfold/lib/openmm.patch && \\\n",
-        "          popd\n",
+        "    %shell conda install -y -q -c conda-forge openmm=7.5.1\n",
+        "    # Apply OpenMM patch.\n",
+        "    %shell pushd /opt/conda/lib/python3.8/site-packages/ && \\\n",
+        "        patch -p0 < /content/openfold/lib/openmm.patch && \\\n",
+        "        popd\n",
+        "    %shell conda install -y -q -c conda-forge pdbfixer=1.7\n",
         "\n",
         "    if(weight_set == 'AlphaFold'):\n",
         "      %shell mkdir --parents \"{ALPHAFOLD_PARAMS_DIR}\"\n",
@@ -222,8 +220,8 @@
         "import unittest.mock\n",
         "import sys\n",
         "\n",
-        "sys.path.insert(0, '/usr/local/lib/python3.7/site-packages/')\n",
-        "sys.path.append('/opt/conda/lib/python3.7/site-packages')\n",
+        "sys.path.insert(0, '/usr/local/lib/python3.8/site-packages/')\n",
+        "sys.path.append('/opt/conda/lib/python3.8/site-packages')\n",
         "\n",
         "# Allows us to skip installing these packages\n",
         "unnecessary_modules = [\n",
@@ -247,6 +245,14 @@
         "import numpy as np\n",
         "import py3Dmol\n",
         "import torch\n",
+        "import shutil\n",
+        "\n",
+        "# Prevent shell magic being broken by openmm, prevent this cryptic error:\n",
+        "# \"NotImplementedError: A UTF-8 locale is required. Got ANSI_X3.4-1968\"\n",
+        "import locale\n",
+        "def getpreferredencoding(do_setlocale = True):\n",
+        "    return \"UTF-8\"\n",
+        "locale.getpreferredencoding = getpreferredencoding\n",
         "\n",
         "# A filthy hack to avoid slow Linear layer initialization\n",
         "import openfold.model.primitives\n",
@@ -267,9 +273,8 @@
         "from openfold.data.tools import jackhmmer\n",
         "from openfold.model import model\n",
         "from openfold.np import protein\n",
-        "if(relax_prediction):\n",
-        "  from openfold.np.relax import relax\n",
-        "  from openfold.np.relax import utils\n",
+        "from openfold.np.relax import relax\n",
+        "from openfold.np.relax.utils import overwrite_b_factors\n",
         "from openfold.utils.import_weights import import_jax_weights_\n",
         "from openfold.utils.tensor_utils import tensor_tree_map\n",
         "\n",
@@ -571,14 +576,13 @@
         "    relaxed_pdb, _, _ = amber_relaxer.process(\n",
         "        prot=unrelaxed_proteins[best_model_name]\n",
         "    )\n",
-        "\n",
-        "    # Write out the prediction\n",
-        "    pred_output_path = os.path.join(output_dir, 'selected_prediction.pdb')\n",
-        "    with open(pred_output_path, 'w') as f:\n",
-        "      f.write(relaxed_pdb)\n",
-        "\n",
         "    best_pdb = relaxed_pdb\n",
         "\n",
+        "  # Write out the prediction\n",
+        "  pred_output_path = os.path.join(output_dir, 'selected_prediction.pdb')\n",
+        "  with open(pred_output_path, 'w') as f:\n",
+        "    f.write(best_pdb)\n",
+        "\n",
         "  pbar.update(n=1)  # Finished AMBER relax.\n",
         "\n",
         "# Construct multiclass b-factors to indicate confidence bands\n",
@@ -590,7 +594,7 @@
         "      banded_b_factors.append(idx)\n",
         "      break\n",
         "banded_b_factors = np.array(banded_b_factors)[:, None] * final_atom_mask\n",
-        "to_visualize_pdb = utils.overwrite_b_factors(best_pdb, banded_b_factors)\n",
+        "to_visualize_pdb = overwrite_b_factors(best_pdb, banded_b_factors)\n",
         "\n",
         "# --- Visualise the prediction & confidence ---\n",
         "show_sidechains = True\n",
@@ -688,7 +692,7 @@
         "\n",
         "\n",
         "# --- Download the predictions ---\n",
-        "!zip -q -r {output_dir}.zip {output_dir}\n",
+        "shutil.make_archive(base_name='prediction', format='zip', root_dir=output_dir)\n",
         "files.download(f'{output_dir}.zip')"
       ],
       "execution_count": null,