plddt and Pae plots

.github/workflows/python-package.yml

@@ -38,7 +38,7 @@ jobs:
         pip install coverage-badge
         coverage-badge -o .blob/coverage.svg -f
     - name: Upload coverage badge
-      uses: actions/upload-artifact@v3
+      uses: actions/upload-artifact@v4
       with:
         name: coverage-badge
         path: coverage.svg

.gitignore

@@ -10,3 +10,5 @@ build/
 .coverage
 
 .vscode/
+
+.history/

MANIFEST.in

@@ -0,0 +1,2 @@
+recursive-include abcfold *.js *.css *.html *.svg *.png
+recursive-include abcfold/plots/pae-viewer-main/src/templates *.tpl

abcfold/abcfold.py

@@ -3,28 +3,48 @@
 import sys
 import tempfile
 from pathlib import Path
+from typing import Dict
 
+from abcfold.abc_script_utils import make_dir, setup_logger
 from abcfold.add_mmseqs_msa import add_msa_to_json
-from abcfold.af3_script_utils import make_dir, setup_logger
 from abcfold.argparse_utils import (alphafold_argparse_util,
                                     boltz_argparse_util, chai_argparse_util,
                                     custom_template_argpase_util,
                                     main_argpase_util, mmseqs2_argparse_util,
                                     prediction_argparse_util)
+from abcfold.plots.pae_plot import create_pae_plots
+from abcfold.plots.plddt_plot import plot_plddt
 from abcfold.processoutput.alphafold3 import AlphafoldOutput
 from abcfold.processoutput.boltz import BoltzOutput
 from abcfold.processoutput.chai import ChaiOutput
 from abcfold.run_alphafold3 import run_alphafold3
 
 logger = setup_logger()
 
+PLOTS_DIR = ".plots"
+
 
 def run(args, config, defaults, config_file):
-    """Run AlphaFold3"""
+    """Run ABCFold
+
+    Args:
+        args (argparse.Namespace): Arguments from the command line
+        config (configparser.SafeConfigParser): Config parser object
+        defaults (dict): Default values from the config file
+        config_file (Path): Path to the config file
+
+
+    Raises:
+        SystemExit: If the database directory or model parameters directory is not found
+
+
+    """
+    outputs = []
 
     args.output_dir = Path(args.output_dir)
 
     make_dir(args.output_dir, overwrite=args.override)
+    make_dir(args.output_dir.joinpath(PLOTS_DIR))
 
     updated_config = False
     if args.model_params != defaults["model_params"]:
@@ -45,9 +65,9 @@ def run(args, config, defaults, config_file):
         sys.exit(1)
 
     with open(args.input_json, "r") as f:
-        af3_json = json.load(f)
+        input_params = json.load(f)
 
-    name = af3_json.get("name")
+    name = input_params.get("name")
     if name is None:
         logger.error("Input JSON must contain a 'name' field")
         sys.exit(1)
@@ -72,14 +92,14 @@ def run(args, config, defaults, config_file):
             else:
                 run_json = Path(args.output_json)
 
-            af3_json = add_msa_to_json(
+            input_params = add_msa_to_json(
                 input_json=input_json,
                 templates=args.templates,
                 num_templates=args.num_templates,
                 custom_template=args.custom_template,
                 custom_template_chain=args.custom_template_chain,
                 target_id=args.target_id,
-                af3_json=af3_json,
+                input_params=input_params,
                 output_json=run_json,
                 to_file=True,
             )
@@ -102,11 +122,13 @@ def run(args, config, defaults, config_file):
                 model_params=args.model_params,
                 database_dir=args.database_dir,
                 number_of_models=args.number_of_models,
+                num_recycles=args.num_recycles,
             )
 
             # Need to find the name of the af3_dir
             af3_out_dir = list(args.output_dir.iterdir())[0]
-            _ = AlphafoldOutput(af3_out_dir, name)
+            ao = AlphafoldOutput(af3_out_dir, input_params, name)
+            outputs.append(ao)
 
         if args.boltz1:
             from abcfold.run_boltz import run_boltz
@@ -115,9 +137,13 @@ def run(args, config, defaults, config_file):
                 input_json=run_json,
                 output_dir=args.output_dir,
                 save_input=args.save_input,
+                number_of_models=args.number_of_models,
+                num_recycles=args.num_recycles,
             )
             bolt_out_dir = list(args.output_dir.glob("boltz_results*"))[0]
-            _ = BoltzOutput(bolt_out_dir, name)
+            bo = BoltzOutput(bolt_out_dir, input_params, name)
+            bo.add_plddt_to_cif()
+            outputs.append(bo)
 
         if args.chai1:
             from abcfold.run_chai1 import run_chai
@@ -128,17 +154,52 @@ def run(args, config, defaults, config_file):
                 output_dir=chai_output_dir,
                 save_input=args.save_input,
                 number_of_models=args.number_of_models,
+                num_recycles=args.num_recycles,
             )
 
-            _ = ChaiOutput(chai_output_dir, name)
+            co = ChaiOutput(chai_output_dir, input_params, name)
+            outputs.append(co)
+
+        plots(outputs, args.output_dir.joinpath(PLOTS_DIR))
+
+
+def plots(outputs: list, output_dir: Path):
+    """
+    Generate plots for the output of the different programs
+
+    Args:
+        outputs (list): List of output objects
+
+    """
+    pathway_plots = create_pae_plots(outputs, output_dir=output_dir)
+    plddt_plot_input: Dict[str, list] = {}
+    for output in outputs:
+        if isinstance(output, AlphafoldOutput):
+            for seed in output.seeds:
+                if "Alphafold3" not in plddt_plot_input:
+                    plddt_plot_input["Alphafold3"] = []
+                plddt_plot_input["Alphafold3"].extend(output.cif_files[seed])
+        elif isinstance(output, BoltzOutput):
+
+            plddt_plot_input["Boltz-1"] = output.cif_files
+        elif isinstance(output, ChaiOutput):
+            plddt_plot_input["Chai-1"] = output.cif_files
+
+    plot_plddt(plddt_plot_input, output_name=output_dir.joinpath("plddt_plot.html"))
+
+    pathway_plots["plddt"] = str(output_dir.joinpath("plddt_plot.html").resolve())
+
+    return pathway_plots
 
 
 def main():
+    """
+    Run AlphaFold3 / Boltz1 / Chai-1
+    """
     import argparse
 
     parser = argparse.ArgumentParser(description="Run AlphaFold3 / Boltz1 / Chai-1")
 
-    # Load defaults from config file
     defaults = {}
     config_file = Path(__file__).parent.joinpath("data", "config.ini")
     config = configparser.SafeConfigParser()

abcfold/add_custom_template.py

@@ -4,7 +4,7 @@
 import logging
 import os
 
-from abcfold.af3_script_utils import get_custom_template
+from abcfold.abc_script_utils import get_custom_template
 from abcfold.argparse_utils import custom_template_argpase_util
 
 logger = logging.getLogger("logger")

abcfold/add_mmseqs_msa.py

@@ -13,7 +13,7 @@
 import requests  # type: ignore
 from tqdm.autonotebook import tqdm
 
-from abcfold.af3_script_utils import (align_and_map,
+from abcfold.abc_script_utils import (align_and_map,
                                       extract_sequence_from_mmcif,
                                       get_custom_template, get_mmcif)
 from abcfold.argparse_utils import (custom_template_argpase_util,
@@ -42,15 +42,15 @@ def add_msa_to_json(
     custom_template,
     custom_template_chain,
     target_id,
-    af3_json=None,
+    input_params=None,
     output_json=None,
     to_file=True,
 ):
-    if af3_json is None:
+    if input_params is None:
         with open(input_json, "r") as f:
-            af3_json = json.load(f)
+            input_params = json.load(f)
 
-    for sequence in af3_json["sequences"]:
+    for sequence in input_params["sequences"]:
         if "protein" in sequence:
             input_sequence = sequence["protein"]["sequence"]
             with tempfile.TemporaryDirectory() as tmpdir:
@@ -77,7 +77,13 @@ def add_msa_to_json(
                     # Can only add templates to protein sequences, so check if there
                     # are multiple protein sequences in the input json
                     if (
-                        len([x for x in af3_json["sequences"] if "protein" in x.keys()])
+                        len(
+                            [
+                                x
+                                for x in input_params["sequences"]
+                                if "protein" in x.keys()
+                            ]
+                        )
                         > 1
                         and not target_id
                     ):
@@ -99,13 +105,13 @@ def add_msa_to_json(
     if to_file:
         if output_json:
             with open(output_json, "w") as f:
-                json.dump(af3_json, f)
+                json.dump(input_params, f)
         else:
             output_json = input_json.replace(".json", "_mmseqs.json")
             with open(output_json, "w") as f:
-                json.dump(af3_json, f)
+                json.dump(input_params, f)
 
-    return af3_json
+    return input_params
 
 
 # Code from https://github.com/sokrypton/ColabFold

abcfold/argparse_utils.py

@@ -43,6 +43,12 @@ def prediction_argparse_util(parser):
         default=5,
         help="Number of models to generate",
     )
+    parser.add_argument(
+        "--num_recycles",
+        type=int,
+        default=10,
+        help="Number of recycles to use during the inference",
+    )
     return parser
 
 
@@ -60,6 +66,7 @@ def boltz_argparse_util(parser):
             help="Save the input json file",
             default=False,
         )
+
     return parser
 
 
@@ -70,16 +77,6 @@ def chai_argparse_util(parser):
         action="store_true",
         help="Run Chai-1",
     )
-    # check if save input is in the parser
-    if "--save_input" not in parser._option_string_actions:
-
-        parser.add_argument(
-            "--save_input",
-            action="store_true",
-            help="Save the input json file",
-            default=False,
-        )
-    # add more arguments here
     return parser
 
 
@@ -104,14 +101,13 @@ def alphafold_argparse_util(parser):
         "-a",
         "--alphafold3",
         action="store_true",
-        help="Run Alphafold",
+        help="Run Alphafold3",
     )
 
     parser.add_argument(
         "--override",
         help="Override the existing output directory, if it exists",
         action="store_true",
     )
-    # add more arguments here
 
     return parser

abcfold/chai1/af3_to_chai.py

@@ -215,14 +215,14 @@ def ccd_to_smiles(self, ccd_id: str):
     def add_ligand(self, seq: dict):
         lig_id = seq["ligand"]["id"]
         ligand_str = ""
-        if "ccdCode" in seq["ligand"]:
+        if "ccdCodes" in seq["ligand"]:
             if isinstance(lig_id, list):
                 for i in seq["ligand"]["id"]:
-                    smile = self.ccd_to_smiles(seq["ligand"]["ccdCode"][0])
+                    smile = self.ccd_to_smiles(seq["ligand"]["ccdCodes"][0])
                     if smile:
                         ligand_str += f">ligand|{i}\n{smile}\n"
             else:
-                smile = self.ccd_to_smiles(seq["ligand"]["ccdCode"])
+                smile = self.ccd_to_smiles(seq["ligand"]["ccdCodes"])
                 if smile:
                     ligand_str = f">ligand|{lig_id}\n{smile}\n"
         if "smiles" in seq["ligand"]:

abcfold/chai1/chai.py

@@ -54,7 +54,7 @@ def run_inference_wrapper(
     seed: int | None = None,
     device: str | None = None,
 ):
-    # print(num_diffn_samples)
+
     result = run_inference(
         fasta_file=fasta_file,
         output_dir=output_dir,

abcfold/plots/pae-viewer-main/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Christoph Elfmann
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.