Merge remote-tracking branch 'sergey/beta' into beta

Author: dohyun-s

AlphaFold2.ipynb

@@ -46,13 +46,14 @@
       "source": [
         "<img src=\"https://raw.githubusercontent.com/sokrypton/ColabFold/main/.github/ColabFold_Marv_Logo_Small.png\" height=\"200\" align=\"right\" style=\"height:240px\">\n",
         "\n",
-        "# ColabFold - v1.6.0-beta\n",
-        "**<font color=\"red\">WARNING</font> this \"beta\" notebook is in active development and may be break without warning.** Use the [\"main\"](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb) notebook for most stable version.\n",
+        "# ColabFold - v1.6.0 - beta\n",
+        "**<font color=\"red\">WARNING</font> this \"beta\" notebook is in active development and may be break without warning.** \n",
+        "Use the [\"main\"](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb) notebook for most stable version.\n",
         "\n",
         "---\n",
         "\n",
         "Easy to use protein structure and complex prediction using [AlphaFold2](https://www.nature.com/articles/s41586-021-03819-2) and [Alphafold2-multimer](https://www.biorxiv.org/content/10.1101/2021.10.04.463034v1). Sequence alignments/templates are generated through [MMseqs2](mmseqs.com) and [HHsearch](https://github.com/soedinglab/hh-suite). For more details, see <a href=\"#Instructions\">bottom</a> of the notebook, checkout the [ColabFold GitHub](https://github.com/sokrypton/ColabFold) and read our manuscript. \n",
-        "Old version: [v1.4](https://colab.research.google.com/github/sokrypton/ColabFold/blob/v1.4.0/AlphaFold2.ipynb)\n",
+        "Old versions: [v1.4](https://colab.research.google.com/github/sokrypton/ColabFold/blob/v1.4.0/AlphaFold2.ipynb), [v1.5](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb)\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",
@@ -226,8 +227,6 @@
         "  # install dependencies\n",
         "  # We have to use \"--no-warn-conflicts\" because colab already has a lot preinstalled with requirements different to ours\n",
         "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold@beta\"\n",
-        "  # high risk high gain\n",
-        "  pip install -q \"jax[cuda11_cudnn805]>=0.3.8,<0.4\" -f https://storage.googleapis.com/jax-releases/jax_releases.html\n",
         "\n",
         "  # for debugging\n",
         "  ln -s /usr/local/lib/python3.*/dist-packages/colabfold colabfold\n",
@@ -266,7 +265,7 @@
       "source": [
         "#@title Run Prediction\n",
         "#@markdown ### Advanced settings\n",
-        "model_type = \"auto\" #@param [\"auto\", \"alphafold2_ptm\", \"alphafold2_multimer_v1\", \"alphafold2_multimer_v2\", \"alphafold2_multimer_v3\"]\n",
+        "model_type = \"auto\" #@param [\"auto\", \"alphafold2\", \"alphafold2_ptm\", \"alphafold2_multimer_v1\", \"alphafold2_multimer_v2\", \"alphafold2_multimer_v3\"]\n",
         "#@markdown - if `auto` will use `alphafold2_ptm` for monomer inputs and `alphafold2_multimer_v3` for multimer inputs.\n",
         "num_recycles = \"auto\" #@param [\"auto\", \"0\", \"1\", \"3\", \"6\", \"12\", \"24\", \"48\"]\n",
         "recycle_early_stop_tolerance = \"auto\" #@param [\"auto\", \"0.0\", \"0.5\", \"1.0\"]\n",
@@ -296,7 +295,7 @@
         "from pathlib import Path\n",
         "from colabfold.download import download_alphafold_params, default_data_dir\n",
         "from colabfold.utils import setup_logging\n",
-        "from colabfold.run_alphafold import get_queries, run, set_model_type\n",
+        "from colabfold.run_alphafold import run, set_model_type\n",
         "from colabfold.inputs import get_queries\n",
         "from colabfold.plot import plot_msa\n",
         "\n",
@@ -307,12 +306,12 @@
         "except:\n",
         "  K80_chk = \"0\"\n",
         "  pass\n",
+        "\n",
         "if \"1\" in K80_chk:\n",
         "  print(\"WARNING: found GPU Tesla K80: limited to total length < 1000\")\n",
-        "  if \"TF_FORCE_UNIFIED_MEMORY\" in os.environ:\n",
-        "    del os.environ[\"TF_FORCE_UNIFIED_MEMORY\"]\n",
-        "  if \"XLA_PYTHON_CLIENT_MEM_FRACTION\" in os.environ:\n",
-        "    del os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"]\n",
+        "else:\n",
+        "  os.environ[\"TF_FORCE_UNIFIED_MEMORY\"] = \"1\"\n",
+        "  os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"4.0\"\n",
         "\n",
         "from colabfold.plot import plot_protein\n",
         "from pathlib import Path\n",
@@ -480,33 +479,40 @@
         "    img = f.read()\n",
         "  return prefix + base64.b64encode(img).decode('utf-8')\n",
         "\n",
-        "pae = image_to_data_url(os.path.join(jobname,f\"{jobname}{jobname_prefix}_pae.png\"))\n",
+        "pae_png_path = os.path.join(jobname,f\"{jobname}{jobname_prefix}_pae.png\")\n",
+        "pae = image_to_data_url(pae_png_path) if os.path.exists(pae_png_path) else None\n",
         "cov = image_to_data_url(os.path.join(jobname,f\"{jobname}{jobname_prefix}_coverage.png\"))\n",
         "plddt = image_to_data_url(os.path.join(jobname,f\"{jobname}{jobname_prefix}_plddt.png\"))\n",
-        "display(HTML(f\"\"\"\n",
-        "<style>\n",
-        "  img {{\n",
-        "    float:left;\n",
-        "  }}\n",
-        "  .full {{\n",
-        "    max-width:100%;\n",
-        "  }}\n",
-        "  .half {{\n",
-        "    max-width:50%;\n",
-        "  }}\n",
-        "  @media (max-width:640px) {{\n",
-        "    .half {{\n",
-        "      max-width:100%;\n",
-        "    }}\n",
-        "  }}\n",
-        "</style>\n",
-        "<div style=\"max-width:90%; padding:2em;\">\n",
-        "  <h1>Plots for {escape(jobname)}</h1>\n",
-        "  <img src=\"{pae}\" class=\"full\" />\n",
-        "  <img src=\"{cov}\" class=\"half\" />\n",
-        "  <img src=\"{plddt}\" class=\"half\" />\n",
-        "</div>\n",
-        "\"\"\"))\n"
+        "\n",
+        "if pae is None:\n",
+        "  display(HTML(f\"\"\"\n",
+        "  <style>\n",
+        "    img {{float:left;}}\n",
+        "    .full {{max-width:100%;}}\n",
+        "    .half {{max-width:50%;}}\n",
+        "    @media (max-width:640px) {{.half {{max-width:100%;}}}}\n",
+        "  </style>\n",
+        "  <div style=\"max-width:90%; padding:2em;\">\n",
+        "    <h1>Plots for {escape(jobname)}</h1>\n",
+        "    <img src=\"{cov}\" class=\"half\" />\n",
+        "    <img src=\"{plddt}\" class=\"half\" />\n",
+        "  </div>\n",
+        "  \"\"\"))\n",
+        "else:\n",
+        "  display(HTML(f\"\"\"\n",
+        "  <style>\n",
+        "    img {{float:left;}}\n",
+        "    .full {{max-width:100%;}}\n",
+        "    .half {{max-width:50%;}}\n",
+        "    @media (max-width:640px) {{.half {{max-width:100%;}}}}\n",
+        "  </style>\n",
+        "  <div style=\"max-width:90%; padding:2em;\">\n",
+        "    <h1>Plots for {escape(jobname)}</h1>\n",
+        "    <img src=\"{pae}\" class=\"full\" />\n",
+        "    <img src=\"{cov}\" class=\"half\" />\n",
+        "    <img src=\"{plddt}\" class=\"half\" />\n",
+        "  </div>\n",
+        "  \"\"\"))"
       ],
       "execution_count": null,
       "outputs": []

README.md

@@ -36,7 +36,7 @@ Link to ["main"](https://github.com/sokrypton/ColabFold) branch.
   - Yes, but be **CAREFUL**, the bfactor column is populated with pLDDT confidence values (higher = better). Phenix.phaser expects a "real" bfactor, where (lower = better). See [post](https://twitter.com/cheshireminima/status/1423929241675120643) from Claudia Millán.
 - What is the maximum length?
   - Limits depends on free GPU provided by Google-Colab `fingers-crossed`
-  - For GPU: `Tesla T4` or `Tesla P100` with ~16G the max length is ~1400
+  - For GPU: `Tesla T4` or `Tesla P100` with ~16G the max length is ~2000
   - For GPU: `Tesla K80` with ~12G the max length is ~1000
   - To check what GPU you got, open a new code cell and type `!nvidia-smi`
 - Is it okay to use the MMseqs2 MSA server (`cf.run_mmseqs2`) on a local computer?
@@ -62,34 +62,8 @@ Link to ["main"](https://github.com/sokrypton/ColabFold) branch.
 - Is there a way to amber-relax structures without having to rerun alphafold/colabfold from scratch?
   - Yes, see this [notebook](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/relax_amber.ipynb).
 
-
 ### Running locally
-
-_Note: Checkout [localcolabfold](https://github.com/YoshitakaMo/localcolabfold) too
-
-It is recommended that you create a conda environment with python version 3.7. If you use a newer python version, you might run into problems when installing tensorflow, as the required version may not be found.
-
-```shell
-conda create --name my_colabfold python=3.7
-```
-Then activate it with `activate my_colabfold`.
-
-Install ColabFold using the `pip` commands below. `pip` will resolve and install all required dependencies and ColabFold should be ready within a few minutes to use. Please check the [JAX documentation](https://github.com/google/jax#pip-installation-gpu-cuda) for how to get JAX to work on your GPU or TPU.
-
-```shell
-pip install "colabfold[alphafold] @ git+https://github.com/sokrypton/ColabFold"
-pip install -q "jax[cuda]>=0.3.8,<0.4" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-# For template-based predictions also install kalign and hhsuite
-conda install -c conda-forge -c bioconda kalign2=2.04 hhsuite=3.3.0
-# For amber also install openmm and pdbfixer
-conda install -c conda-forge openmm=7.5.1 pdbfixer
-```
-
-```shell
-colabfold_batch <directory_with_fasta_files> <result_dir> 
-```
-
-If no GPU or TPU is present, `colabfold_batch` can be executed (slowly) using only a CPU with the `--cpu` parameter.
+For instructions on how to install ColabFold locally see: [localcolabfold](https://github.com/YoshitakaMo/localcolabfold) 
 
 ### Generating MSAs for large scale structure/complex predictions
 
@@ -153,15 +127,15 @@ Searches against the ColabFoldDB can be done in two different modes:
 -----------------
 **OLD Updates**
 ```diff
-  2023/01/03: The MSA server's faulty hardware from 12/26 was replaced.
-              There were intermittent failures on 12/26 and 1/3. Currently, 
-              there are no known issues. Let us know if you experience any.
-  2022/10/10: Bugfix: random_seed was not being used for alphafold-multimer.
-              Same structure was returned regardless of defined seed. This
-              has been fixed!
-  2022/07/13: We have set up a new ColabFold MSA server provided by Korean
-              Bioinformation Center. It provides accelerated MSA generation, 
-              we updated the UniRef30 to 2022_02 and PDB/PDB70 to 220313.
+  03Jan2023: The MSA server's faulty hardware from 12/26 was replaced.
+             There were intermittent failures on 12/26 and 1/3. Currently, 
+             there are no known issues. Let us know if you experience any.
+  10Oct2022: Bugfix: random_seed was not being used for alphafold-multimer.
+             Same structure was returned regardless of defined seed. This
+             has been fixed!
+  13Jul2022: We have set up a new ColabFold MSA server provided by Korean
+             Bioinformation Center. It provides accelerated MSA generation, 
+             we updated the UniRef30 to 2022_02 and PDB/PDB70 to 220313.
   11Mar2022: We use in default AlphaFold-multimer-v2 weights for complex modeling. 
              We also offer the old complex modes "AlphaFold-ptm" or "AlphaFold-multimer-v1"
   04Mar2022: ColabFold now uses a much more powerful server for MSAs and searches through the ColabFoldDB instead of BFD/MGnify. 

colabfold/alphafold/models.py

@@ -8,6 +8,7 @@
 
 
 def load_models_and_params(
+  num_models: int,
   use_templates: bool,
   num_recycles: Optional[int] = None,
   recycle_early_stop_tolerance: Optional[float] = None,
@@ -22,7 +23,7 @@ def load_models_and_params(
   use_fuse: bool = True,
   use_bfloat16: bool = True,
   use_dropout: bool = False,
-
+  save_all: bool = False,
 ) -> List[Tuple[str, model.RunModel, haiku.Params]]:
   """We use only two actual models and swap the parameters to avoid recompiling.
 
@@ -77,6 +78,12 @@ def load_models_and_params(
         else:
           model_config.data.common.max_extra_msa = max_extra_seq
 
+      # disable some outputs if not being saved
+      if not save_all:
+        model_config.model.heads.distogram.weight = 0.0
+        model_config.model.heads.masked_msa.weight = 0.0
+        model_config.model.heads.experimentally_resolved.weight = 0.0
+
       # set number of recycles and ensembles      
       if "multimer" in model_suffix:
         if num_recycles is not None:
@@ -117,6 +124,8 @@ def load_models_and_params(
     )
   # reorder model
   for n, model_number in enumerate(model_order):
+    if n == num_models:
+      break
     model_name = f"model_{model_number}"
     for m in model_runner_and_params_build_order:
       if model_name == m[0]:

colabfold/alphafold/relax.py

@@ -1,7 +1,7 @@
 #############
 # relax functions
 #############
-
+from pathlib import Path
 from alphafold.relax import relax
 from alphafold.common import protein