v2.3.1 - merge 'beta' into 'main' (#371)

* Added input_features_callback to run() function to allow plotting MSA before structure predictions computation.

* Added lower level plot_pae() and plot_protein_backbone() functions to allow making individual plots instead of multi-pane composite plots.

* Fix typo in new protein backbone plotting code.

* testing v.2.3.0

* adding support for multimer_v3

* updating version number of alphafold-colabfold

* Update batch.py

convert unrelaxed protein features to numpy before saving (to speedup save function)

* add option to enable/disable fuse

* adding option to enable/disable fuse

* download v3 params by default

* adding option for iterating through random_seed(s)

* adding option to iterate through random seeds

* Update poetry.lock

* Delete poetry.lock

* Update batch.py

* Update test.yml

only run test on main branch

* fixing encoding error (when amber relax is used and notebook is run multiple times)

* refactoring code a bit to write the results to jobname directory

* add model_type into filename

* adding model_type to filename

* disable bfloat16 for old (v1, v2) multimer models

* select best view using plddt

* adding option to "save all" outputs of model

* adding option to save_all

* adding max_msa option to notebook

* allow max_msa settings for multimer

* adjust max settings

* updating dm-haiku to fix error in jax

* adding option to disable use_cluster_profile

* Created using Colaboratory

* code cleanup

* code cleanup

* Update models.py

* add option to enable bfloat16

* Fix tensorflow import

* making bfloat16 and fuse the default

* cleanup

* major bugfix for ptm + is_complex (#360)

* Update README.md

* adding option to control num of top models to relax

* control number of models to relax

* Update batch.py

* update pae filename

* remove whitespace from saved images

* addressing memory leak and pae bugfix

all outputs are saved as they are generated to avoid memory leaks. scores reloaded later for plotting.

bugfix: *_predicted_aligned_error_v1.json is now the best ranked output (before it was whichever model was saved first)

* adding leading zeros to rank id and seed id

* adding leading zeros to rank and seed id(s)

* Delete colabfold_alphafold.py

unused file

* monomers can now be predicted with multimer model!

* fix numbering

* Update batch.py

bugfix for custom msa input
run function updated to return model ranks and metrics

* update run function to print metrics and keep track of best model rank

* replacing max_msa

replacing max_msa option with max_seq and max_extra_seq to make easier to control.

* Update batch.py

* Update batch.py

rank_by wasn't computed correctly for multimer ranking

* Update AlphaFold2_batch.ipynb

* Update AlphaFold2_batch.ipynb

* adding TPU support

* remove TPU warning...

* disable BiopythonDeprecationWarning

* Move load_models_and_params to before the job loop (#368)

* Update batch.py

Move load_models_and_params to before the job loop to avoid issues when resuming from an existing directory with already finished jobs.

* Update batch.py

Change to check if it's the first job instead of moving load_models_and_params

* recompile_padding update

padding is now defined as a constant int instead of float

* Update README.md

* Created using Colaboratory

* Update batch.py

* Created using Colaboratory

* adding iptm support to ptm (for complexes)

* removing extra print out

* Update batch.py

* Update batch.py

* add option to save intermediate results

* add option to save_recycles

* Update batch.py

* updating pyproject.toml to use latest alphafold-colabfold v2.3

* Revert "updating pyproject.toml to use latest alphafold-colabfold v2.3"

This reverts commit 05ec548.

* Update alphafold-colabfold to 2.3.1

* attempt to fix test

* Adding link back to old version

* update test data

* Update test_colabfold.py

* fix test (fingers crossed)

* Update test_colabfold.py

* Update README.md

* updating notebooks to use "main" instead of "beta"

---------

Co-authored-by: Tom Goddard <[email protected]>
Co-authored-by: YoshitakaMo <[email protected]>
Co-authored-by: Milot Mirdita <[email protected]>
Co-authored-by: Dennis Svedberg <[email protected]>
Co-authored-by: Martin Steinegger <[email protected]>
Co-authored-by: Martin Steinegger <[email protected]>

Author: 6 people

.github/workflows/test.yml

@@ -1,7 +1,12 @@
 name: Test
 
-on: [ push, pull_request ]
-
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
 jobs:
   run-tests:
     runs-on: ubuntu-latest

.gitignore

@@ -7,3 +7,4 @@
 /tmalign
 /pytestemp
 __pycache__
+.DS_Store

AlphaFold2.ipynb

@@ -1,15 +1,7 @@
-# ColabFold
+# ColabFold - v1.5.0
 
 ```diff
-+ 2022/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.
++ 04Feb2023: ColabFold updated to use AlphaFold v2.3.1!
 ```
 <p align="center"><img src="https://github.com/sokrypton/ColabFold/raw/main/.github/ColabFold_Marv_Logo.png" height="250"/></p>
 
@@ -28,7 +20,7 @@
 | [OmegaFold](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/omegafold.ipynb) | Yes | Maybe | No | No| No |
 ||
 | **OLD retired notebooks** | | | | | |
-| [AlphaFold2_complexes](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2_complexes.ipynb) | No | Yes | No | No | No | 
+| [AlphaFold2_complexes](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/AlphaFold2_complexes.ipynb) | No | Yes | No | No | No | 
 | [AlphaFold2_jackhmmer](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/AlphaFold_wJackhmmer.ipynb) | Yes | No | Yes | Yes | No |
 | [AlphaFold2_noTemplates_noMD](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/verbose/alphafold_noTemplates_noMD.ipynb) |
 | [AlphaFold2_noTemplates_yesMD](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/verbose/alphafold_noTemplates_yesMD.ipynb) |
@@ -158,6 +150,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.
   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. 
@@ -196,4 +197,3 @@ Searches against the ColabFoldDB can be done in two different modes:
   20Nov2021  "AMBER" is fixed thanks to Kevin Pan
 ```
 -----------------
-

README.md

@@ -260,4 +260,4 @@
       ]
     }
   ]
-}
+}

batch/AlphaFold2_batch.ipynb

@@ -24,7 +24,7 @@
         "colab_type": "text"
       },
       "source": [
-        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/AlphaFold2_advanced.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/beta/beta/AlphaFold2_advanced.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
       ]
     },
     {

beta/AlphaFold2_advanced.ipynb

@@ -1,9 +1,7 @@
 from pathlib import Path
 from functools import wraps, partialmethod
 from typing import Tuple, List, Optional
-
 import haiku
-
 from alphafold.model import model, config, data
 from alphafold.model.modules import AlphaFold
 from alphafold.model.modules_multimer import AlphaFold as AlphaFoldMultimer
@@ -12,129 +10,121 @@
 def load_models_and_params(
     num_models: int,
     use_templates: bool,
-    num_recycle: int = 3,
+    num_recycles: Optional[int] = None,
+    recycle_early_stop_tolerance: Optional[float] = None,
     num_ensemble: int = 1,
     model_order: Optional[List[int]] = None,
     model_suffix: str = "_ptm",
     data_dir: Path = Path("."),
-    recompile_all_models: bool = False,
     stop_at_score: float = 100,
-    rank_by: str = "plddt",
-    return_representations: bool = False,
-    training: bool = False,
-    max_msa: str = None,
+    rank_by: str = "auto",
+    max_seq: Optional[int] = None,
+    max_extra_seq: Optional[int] = None,
+    use_cluster_profile: Optional[bool] = None,
+    use_fuse: bool = True,
+    use_bfloat16: bool = True,
+    use_dropout: bool = False,
+
 ) -> List[Tuple[str, model.RunModel, haiku.Params]]:
     """We use only two actual models and swap the parameters to avoid recompiling.
 
     Note that models 1 and 2 have a different number of parameters compared to models 3, 4 and 5,
     so we load model 1 and model 3.
     """
 
-    if return_representations:
-        # this forces the AlphaFold to always return representations
-        AlphaFold.__call__ = partialmethod(
-            AlphaFold.__call__, return_representations=True
-        )
-
-        AlphaFoldMultimer.__call__ = partialmethod(
-            AlphaFoldMultimer.__call__, return_representations=True
-        )
-
-    if not model_order:
-        model_order = [3, 4, 5, 1, 2]
-
     # Use only two model and later swap params to avoid recompiling
     model_runner_and_params: [Tuple[str, model.RunModel, haiku.Params]] = []
 
-    if recompile_all_models:
-        for n, model_number in enumerate(model_order):
-            if n == num_models:
-                break
-            model_name = f"model_{model_number}"
-            params = data.get_model_haiku_params(
-                model_name=model_name + model_suffix, data_dir=str(data_dir)
-            )
-            model_config = config.model_config(model_name + model_suffix)
-            model_config.model.stop_at_score = float(stop_at_score)
-            model_config.model.stop_at_score_ranker = rank_by
-            if max_msa != None:
-                max_msa_clusters, max_extra_msa = [int(x) for x in max_msa.split(":")]
-                model_config.data.eval.max_msa_clusters = max_msa_clusters
-                model_config.data.common.max_extra_msa = max_extra_msa
-            if model_suffix == "_ptm":
-                model_config.data.common.num_recycle = num_recycle
-                model_config.model.num_recycle = num_recycle
-                model_config.data.eval.num_ensemble = num_ensemble
-            elif model_suffix.startswith("_multimer"):
-                model_config.model.num_recycle = num_recycle
-                if training:
-                    model_config.model.num_ensemble_train = num_ensemble
-                else:
-                    model_config.model.num_ensemble_eval = num_ensemble
-            model_runner_and_params.append(
-                (
-                    model_name,
-                    model.RunModel(model_config, params, is_training=training),
-                    params,
-                )
-            )
+    if model_order is None: model_order = [1, 2, 3, 4, 5]
+
+    model_build_order = [3, 4, 5, 1, 2]
+    if "multimer" in model_suffix:
+        models_need_compilation = [3]
     else:
+        # only models 1,2 use templates
         models_need_compilation = [1, 3] if use_templates else [3]
-        model_build_order = [3, 4, 5, 1, 2]
-        model_runner_and_params_build_order: [
-            Tuple[str, model.RunModel, haiku.Params]
-        ] = []
-        model_runner = None
-        for model_number in model_build_order:
-            if model_number in models_need_compilation:
-                model_config = config.model_config(
-                    "model_" + str(model_number) + model_suffix
-                )
-                model_config.model.stop_at_score = float(stop_at_score)
-                model_config.model.stop_at_score_ranker = rank_by
-                if max_msa != None:
-                    max_msa_clusters, max_extra_msa = [
-                        int(x) for x in max_msa.split(":")
-                    ]
-                    model_config.data.eval.max_msa_clusters = max_msa_clusters
-                    model_config.data.common.max_extra_msa = max_extra_msa
-                if model_suffix == "_ptm":
-                    model_config.data.common.num_recycle = num_recycle
-                    model_config.model.num_recycle = num_recycle
-                    model_config.data.eval.num_ensemble = num_ensemble
-                elif model_suffix.startswith("_multimer"):
-                    model_config.model.num_recycle = num_recycle
-                    if training:
-                        model_config.model.num_ensemble_train = num_ensemble
-                    else:
-                        model_config.model.num_ensemble_eval = num_ensemble
-                model_runner = model.RunModel(
-                    model_config,
-                    data.get_model_haiku_params(
-                        model_name="model_" + str(model_number) + model_suffix,
-                        data_dir=str(data_dir),
-                    ),
-                    is_training=training,
-                )
-            model_name = f"model_{model_number}"
+    
+    model_runner_and_params_build_order: [Tuple[str, model.RunModel, haiku.Params]] = []
+    model_runner = None
+    for model_number in model_build_order:
+        if model_number in models_need_compilation:
+
+            # get configurations
+            model_config = config.model_config("model_" + str(model_number) + model_suffix)
+            model_config.model.stop_at_score = float(stop_at_score)
+            model_config.model.rank_by = rank_by
+
+            # set dropouts
+            model_config.model.global_config.eval_dropout = use_dropout
+
+            # set bfloat options
+            model_config.model.global_config.bfloat16 = use_bfloat16
+            
+            # set fuse options
+            model_config.model.embeddings_and_evoformer.evoformer.triangle_multiplication_incoming.fuse_projection_weights = use_fuse
+            model_config.model.embeddings_and_evoformer.evoformer.triangle_multiplication_outgoing.fuse_projection_weights = use_fuse
+            if "multimer" in model_suffix or model_number in [1,2]:
+                model_config.model.embeddings_and_evoformer.template.template_pair_stack.triangle_multiplication_incoming.fuse_projection_weights = use_fuse
+                model_config.model.embeddings_and_evoformer.template.template_pair_stack.triangle_multiplication_outgoing.fuse_projection_weights = use_fuse
+                        
+            # set number of sequences options
+            if max_seq is not None:
+                if "multimer" in model_suffix:
+                    model_config.model.embeddings_and_evoformer.num_msa = max_seq
+                else:
+                    model_config.data.eval.max_msa_clusters = max_seq
+            
+            if max_extra_seq is not None:
+                if "multimer" in model_suffix:
+                    model_config.model.embeddings_and_evoformer.num_extra_msa = max_extra_seq
+                else:
+                    model_config.data.common.max_extra_msa = max_extra_seq
+
+            # set number of recycles and ensembles            
+            if "multimer" in model_suffix:
+                if num_recycles is not None:
+                    model_config.model.num_recycle = num_recycles
+                if use_cluster_profile is not None:
+                    model_config.model.embeddings_and_evoformer.use_cluster_profile = use_cluster_profile
+                model_config.model.num_ensemble_eval = num_ensemble
+            else:
+                if num_recycles is not None:
+                    model_config.data.common.num_recycle = num_recycles
+                    model_config.model.num_recycle = num_recycles
+                model_config.data.eval.num_ensemble = num_ensemble
+
+
+            if recycle_early_stop_tolerance is not None:
+                model_config.model.recycle_early_stop_tolerance = recycle_early_stop_tolerance
+            
+            # get model runner
             params = data.get_model_haiku_params(
-                model_name=model_name + model_suffix, data_dir=str(data_dir)
+                model_name="model_" + str(model_number) + model_suffix,
+                data_dir=str(data_dir), fuse=use_fuse)
+            model_runner = model.RunModel(
+                model_config,
+                params,
             )
-            # keep only parameters of compiled model
-            params_subset = {}
-            for k in model_runner.params.keys():
-                params_subset[k] = params[k]
+        
+        model_name = f"model_{model_number}"
+        params = data.get_model_haiku_params(
+            model_name=model_name + model_suffix, data_dir=str(data_dir), fuse=use_fuse,
+        )
+        # keep only parameters of compiled model
+        params_subset = {}
+        for k in model_runner.params.keys():
+            params_subset[k] = params[k]
 
-            model_runner_and_params_build_order.append(
-                (model_name, model_runner, params_subset)
-            )
-        # reorder model
-        for n, model_number in enumerate(model_order):
-            if n == num_models:
+        model_runner_and_params_build_order.append(
+            (model_name, model_runner, params_subset)
+        )
+    # 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]:
+                model_runner_and_params.append(m)
                 break
-            model_name = f"model_{model_number}"
-            for m in model_runner_and_params_build_order:
-                if model_name == m[0]:
-                    model_runner_and_params.append(m)
-                    break
     return model_runner_and_params