Merge pull request #308 from vaclavhanzl/colab-add-dep-fix-python-ver…

…sion Fix Colab again after ModelCIF merge and python version change in Colab
aqlaboratory · May 17, 2023 · 1f9270c · 1f9270c
2 parents cc1f8ba + 35b1c4e
commit 1f9270c
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 14 deletions.
diff --git a/notebooks/OpenFold.ipynb b/notebooks/OpenFold.ipynb
@@ -97,13 +97,16 @@
         "#@markdown **Note**: This installs the software on the Colab \n",
         "#@markdown notebook in the cloud and not on your computer.\n",
         "\n",
+        "import sys\n",
         "from IPython.utils import io\n",
         "import os\n",
         "import subprocess\n",
         "import tqdm.notebook\n",
         "\n",
         "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
         "\n",
+        "python_version = '.'.join(sys.version.split('.')[:2]) #get string like \"3.9\"\n",
+        "\n",
         "try:\n",
         "  with io.capture_output() as captured:\n",
         "    %shell sudo apt install --quiet --yes hmmer\n",
@@ -125,12 +128,13 @@
         "    %shell conda install -y -q -c conda-forge -c bioconda \\\n",
         "      kalign2=2.04 \\\n",
         "      hhsuite=3.3.0 \\\n",
-        "      python=3.8 \\\n",
+        "      python={python_version} \\\n",
         "      2>&1 1>/dev/null\n",
         "    %shell pip install -q \\\n",
         "      ml-collections==0.1.0 \\\n",
         "      PyYAML==5.4.1 \\\n",
-        "      biopython==1.79\n",
+        "      biopython==1.79 \\\n",
+        "      modelcif==0.7\n",
         "\n",
         "    # Create a ramdisk to store a database chunk to make Jackhmmer run fast.\n",
         "    %shell sudo mkdir -m 777 --parents /tmp/ramdisk\n",
@@ -181,12 +185,8 @@
         "    %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",
-        "    %shell conda install -y -q -c conda-forge pdbfixer=1.7\n",
+        "    %shell conda install -y -q -c conda-forge openmm=7.7.0\n",
+        "    %shell conda install -y -q -c conda-forge pdbfixer\n",
         "\n",
         "    if(weight_set == 'AlphaFold'):\n",
         "      %shell mkdir --parents \"{ALPHAFOLD_PARAMS_DIR}\"\n",
@@ -220,8 +220,8 @@
         "import unittest.mock\n",
         "import sys\n",
         "\n",
-        "sys.path.insert(0, '/usr/local/lib/python3.8/site-packages/')\n",
-        "sys.path.append('/opt/conda/lib/python3.8/site-packages')\n",
+        "sys.path.insert(0, f'/usr/local/lib/python{python_version}/site-packages/')\n",
+        "sys.path.append(f'/opt/conda/lib/python{python_version}/site-packages')\n",
         "\n",
         "# Allows us to skip installing these packages\n",
         "unnecessary_modules = [\n",
@@ -791,4 +791,4 @@
       ]
     }
   ]
-}
+}
diff --git a/openfold/np/relax/relax.py b/openfold/np/relax/relax.py
@@ -57,7 +57,7 @@ def __init__(
         self._use_gpu = use_gpu
 
     def process(
-        self, *, prot: protein.Protein, cif_output: bool
+        self, *, prot: protein.Protein, cif_output: bool = False
     ) -> Tuple[str, Dict[str, Any], np.ndarray]:
         """Runs Amber relax on a prediction, adds hydrogens, returns PDB string."""
         out = amber_minimize.run_pipeline(

diff --git a/openfold/utils/script_utils.py b/openfold/utils/script_utils.py
@@ -228,7 +228,7 @@ def prep_output(out, batch, feature_dict, feature_processor, config_preset, mult
     return unrelaxed_protein
 
 
-def relax_protein(config, model_device, unrelaxed_protein, output_directory, output_name, cif_output):
+def relax_protein(config, model_device, unrelaxed_protein, output_directory, output_name, cif_output=False):
     amber_relaxer = relax.AmberRelaxation(
         use_gpu=(model_device != "cpu"),
         **config.relax,
@@ -257,4 +257,4 @@ def relax_protein(config, model_device, unrelaxed_protein, output_directory, out
     with open(relaxed_output_path, 'w') as fp:
         fp.write(struct_str)
 
-    logger.info(f"Relaxed output written to {relaxed_output_path}...")
+    logger.info(f"Relaxed output written to {relaxed_output_path}...")