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

* Python 3.8 is now what we get in Colab, use it in our Conda as well
* Fix imports, make some unconditional - chain of depencencies makes them
  needed even for relax_prediction==False
* Mitigate spurious error "A UTF-8 locale is required. Got ANSI_X3.4-1968"
  which makes all shell commands defunct - use shutil.make_archive()

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,14 @@
         "    %shell cp -f /content/stereo_chemical_props.txt /content/openfold/openfold/resources\n",
         "    %shell /usr/bin/python3 -m pip install -q ./openfold\n",
         "\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",
+        "\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 pdbfixer=1.7\n",
         "\n",
         "    if(weight_set == 'AlphaFold'):\n",
         "      %shell mkdir --parents \"{ALPHAFOLD_PARAMS_DIR}\"\n",
@@ -222,8 +222,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 +247,7 @@
         "import numpy as np\n",
         "import py3Dmol\n",
         "import torch\n",
+        "import shutil\n",
         "\n",
         "# A filthy hack to avoid slow Linear layer initialization\n",
         "import openfold.model.primitives\n",
@@ -269,7 +270,7 @@
         "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.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",
@@ -590,7 +591,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 +689,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,