Custom template fix, IPSAE, and Protenix

README.md

@@ -4,7 +4,7 @@
 ![Coverage](https://raw.githubusercontent.com/rigdenlab/ABCFold/refs/heads/main/.blob/coverage.svg)
 
 
-Scripts to run AlphaFold3, Boltz and Chai-1 with MMseqs2 Multiple sequence alignments (MSAs) and custom templates.
+Scripts to run AlphaFold3, Boltz, Chai-1 and Protenix with MMseqs2 Multiple sequence alignments (MSAs) and custom templates.
 
 ## Table of Contents
 - [Installation](#installation)
@@ -14,21 +14,21 @@ Scripts to run AlphaFold3, Boltz and Chai-1 with MMseqs2 Multiple sequence align
 
 ## Installation
 
-We recommend installing this package in a virtual environment or conda / micromamba environment. Python 3.11 is recommended, but the package should work with Python 3.9 and above.
+We recommend installing this package in a micromamba environment. Python 3.11 is recommended, but the package should work with Python 3.9 and above.
+For instructions on installing micromamba, check their website [here](https://mamba.readthedocs.io/en/latest/installation/micromamba-installation.html).
 
-To set up a conda/micromamba environment, run:
+To set up a micromamba environment, run:
 ```bash
-conda create -n abcfold python=3.11
-conda activate abcfold
-```
-
-or
-
-```bash
-micromamba env create -n abcfold python=3.11
+micromamba env create -n abcfold python=3.11 -y
 micromamba activate abcfold
 ```
 
+> **Environment note**
+>
+> - Your main Python environment can be Conda, micromamba, or virtualenv.
+> - `abcfold` will automatically create **internal micromamba environments** to run Boltz, Chai-1 and Protenix safely, therefore a micromamba installation is required.
+> - This prevents package conflicts with your main environment and ensures reproducible results.
+
 
 To install the package from PyPI, run:
 
@@ -57,7 +57,7 @@ python -m pre_commit install
 
 ### Running ABCfold
 
-ABCFold will run Alphafold3, Boltz and Chai-1 consecutively. The program takes an input of a JSON in the Alphafold3 format (For full instruction on how to format this, click [here](https://github.com/google-deepmind/alphafold3/blob/main/docs/input.md)). An example JSON is shown below:
+ABCFold will run Alphafold3, Boltz, Chai-1 and Protenix consecutively. The program takes an input of a JSON in the Alphafold3 format (For full instruction on how to format this, click [here](https://github.com/google-deepmind/alphafold3/blob/main/docs/input.md)). An example JSON is shown below:
 
 ```json
 {
@@ -80,7 +80,7 @@ Please make sure you have AlphaFold3 installed on your system (Instructions [her
 
 For the majority of jobs, ABCFold can be run as follows:
 ```bash
-abcfold <input_json>  <output_dir> -abc --mmseqs2 --model_params <path_to_af3_model_params>
+abcfold <input_json>  <output_dir> -abcp --mmseqs2 --model_params <path_to_af3_model_params>
 ```
 > [!NOTE]
 > `--model_params` is stored after the first run, therefore subsequent ABCFold jobs don't require this flag.
@@ -100,7 +100,7 @@ However, there you may wish to use the following flags to add run time options s
 #### Main arguments
 - `<input_json>`: Path to the input AlphaFold3 JSON file.
 - `<output_dir>`: Path to the output directory.
-- `-a`, `-b`, `-c` (`--alphafold3`, `--boltz`,`--chai1`): Flags to run Alphafold3, Boltz and Chai-1 respectively. If none of these flags are provided, Alphafold3 will be run by default.
+- `-a`, `-b`, `-c`, `-p` (`--alphafold3`, `--boltz`,`--chai1`,`--protenix`): Flags to run Alphafold3, Boltz, Chai-1 and Protenix respectively. If none of these flags are provided, Alphafold3 will be run by default.
 - `--mmseqs2`: [optional] Flag to use MMseqs2 MSAs and templates (if specified).
 - `--mmseqs_database`: [optional] The path to the database used by a local copy of MMSeqs2, provided mmseqs is installed, the inclusion of this flag allows MMseqs2 to be run locally.
 - `--override`: [optional] Flag to override the existing output directory.
@@ -111,7 +111,7 @@ However, there you may wish to use the following flags to add run time options s
 - `--model_params`: Path to the directory containing the AlphaFold3 model parameters.
 - `--database`: [optional] Path to the directory containing the AlphaFold3 databases #Note: This is not used if using the
 `--mmseqs2` flag.
-- `--sif_path`: [optional] Path to sif file if using an AlphaFold3 singularity instead of Docker
+- `--af3_sif_path`: [optional] Path to sif file if using an AlphaFold3 singularity instead of Docker
 - `--use_af3_template_search`[optional] If providing your own custom MSA or you've ran `--mmseqs2`, allow Alphafold3 to search for templates
 
 #### Template arguments
@@ -132,14 +132,14 @@ However, there you may wish to use the following flags to add run time options s
 If you wanted to provide a custom template, `custom_a.pdb` for your protein sequence with the ID `A` and you have your template has two chains: chain `A` and chain `B` and chain `B` is what you want the template to be, you could run:
 
 ```bash
-abcfold <input_json>  <output_dir> -abc --mmseqs2 --custom_template custom_a.pdb  --custom_template_chain B --target_id A
+abcfold <input_json>  <output_dir> -abcp --mmseqs2 --custom_template custom_a.pdb  --custom_template_chain B --target_id A
 
 ```
 
 If you had multiple IDs in your input sequence, multiple template files and you wanted to provide 3 custom templates, chain `A` from `custom_a.pdb`, chain `B` from `custom_b.pdb`, and chain B from `custom_c.pdb`, where `custom_a.pdb` and `custom_b.pdb` correspond to the ID `A` and `custom_c.pdb` corresponds to the ID `B`, you could run:
 
 ```bash
-abcfold <input_json>  <output_dir> -abc --mmseqs2 --custom_template custom_a.pdb custom_b.pdb custom_c.pdb --custom_template_chain A B B --target_id A A B
+abcfold <input_json>  <output_dir> -abcp --mmseqs2 --custom_template custom_a.pdb custom_b.pdb custom_c.pdb --custom_template_chain A B B --target_id A A B
 
 ```
 ### Output
@@ -168,7 +168,7 @@ you will find `open_output.py` in your `<output_dir>`. This needs to be run from
 Below are scripts for adding MMseqs2 MSAs and custom templates to AlphaFold3 input JSON files.
 
 > [!WARNING]
-> These scripts will only modify the input JSON files, I.E. they will NOT run AlphaFold3, Boltz and Chai-1.
+> These scripts will only modify the input JSON files, I.E. they will NOT run AlphaFold3, Boltz, Chai-1 and Protenix.
 
 ### Adding MMseqs2 MSAs and templates
 

abcfold/abcfold.py

@@ -15,6 +15,7 @@
                                     custom_template_argpase_util,
                                     main_argpase_util, mmseqs2_argparse_util,
                                     prediction_argparse_util,
+                                    protenix_argparse_util,
                                     raise_argument_errors,
                                     visuals_argparse_util)
 from abcfold.html.html_utils import (PORT, NoCacheHTTPRequestHandler,
@@ -26,10 +27,12 @@
 from abcfold.output.boltz import BoltzOutput
 from abcfold.output.chai import ChaiOutput
 from abcfold.output.file_handlers import superpose_models
+from abcfold.output.protenix import ProtenixOutput
 from abcfold.output.utils import (get_gap_indicies, insert_none_by_minus_one,
                                   make_dummy_m8_file)
 from abcfold.scripts.abc_script_utils import (check_input_json, make_dir,
                                               make_dummy_af3_db, setup_logger)
+from abcfold.scripts.add_custom_template import add_custom_template
 from abcfold.scripts.add_mmseqs_msa import add_msa_to_json
 
 logger = setup_logger()
@@ -71,6 +74,10 @@ def run(args, config, defaults, config_file):
     if args.database_dir is not None and args.database_dir != defaults["database_dir"]:
         config.set("Databases", "database_dir", args.database_dir)
         updated_config = True
+    af3_sif = args.af3_sif_path
+    if af3_sif is not None and af3_sif != defaults.get("af3_sif_path"):
+        config.set("Sif_paths", "af3_sif_path", af3_sif)
+        updated_config = True
     if updated_config:
         with open(config_file, "w") as f:
             config.write(f)
@@ -91,21 +98,6 @@ def run(args, config, defaults, config_file):
         logger.error("Input JSON must contain a 'name' field")
         sys.exit(1)
 
-    if args.alphafold3:
-        from abcfold.alphafold3.check_install import check_af3_install
-
-        check_af3_install(interactive=False, sif_path=args.sif_path)
-
-    if args.boltz:
-        from abcfold.boltz.check_install import check_boltz
-
-        check_boltz()
-
-    if args.chai1:
-        from abcfold.chai1.check_install import check_chai1
-
-        check_chai1()
-
     with tempfile.TemporaryDirectory() as temp_dir_str:
         temp_dir = Path(temp_dir_str)
         if args.mmseqs2:
@@ -115,6 +107,7 @@ def run(args, config, defaults, config_file):
                     input_json.name.replace(".json", "_mmseqs.json")
                 )
             else:
+                input_json = Path(args.input_json)
                 run_json = Path(args.output_json)
 
             input_params = add_msa_to_json(
@@ -123,17 +116,23 @@ def run(args, config, defaults, config_file):
                 templates=args.templates,
                 num_templates=args.num_templates,
                 chai_template_output=temp_dir.joinpath("all_chains.m8"),
-                custom_template=args.custom_template,
-                custom_template_chain=args.custom_template_chain,
-                target_id=args.target_id,
                 input_params=input_params,
                 output_json=run_json,
                 to_file=True,
             )
-
         else:
             run_json = Path(args.input_json)
 
+        if args.custom_template is not None:
+            add_custom_template(
+                run_json,
+                args.target_id,
+                args.custom_template,
+                args.custom_template_chain,
+                output_json=run_json,
+                to_file=True
+            )
+
         successful_runs = []
         if args.alphafold3:
             af3_database = args.database_dir
@@ -153,7 +152,7 @@ def run(args, config, defaults, config_file):
                 database_dir=af3_database,
                 number_of_models=args.number_of_models,
                 num_recycles=args.num_recycles,
-                sif_path=args.sif_path,
+                sif_path=af3_sif,
                 save_distogram=args.save_distogram
             )
 
@@ -212,6 +211,25 @@ def run(args, config, defaults, config_file):
                 outputs.append(co)
             successful_runs.append(chai_success)
 
+        if args.protenix:
+            from abcfold.protenix.run_protenix import run_protenix
+
+            protenix_success = run_protenix(
+                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,
+            )
+
+            if protenix_success:
+                protenix_output_dirs = list(args.output_dir.glob("protenix_results*"))
+                po = ProtenixOutput(
+                    protenix_output_dirs, input_params, name, args.save_input
+                )
+                outputs.append(po)
+            successful_runs.append(protenix_success)
+
         if args.no_visuals:
             logger.info("Visuals disabled")
             return
@@ -243,6 +261,7 @@ def run(args, config, defaults, config_file):
                         model.check_clashes()
                         score_file = ao.output[seed][idx]["summary"]
                         plddt = model.residue_plddts
+                        pae = ao.output[seed][idx]["af3_pae"]
                         if len(indicies) > 0:
                             plddt = insert_none_by_minus_one(
                                 indicies[index_counter], plddt
@@ -253,6 +272,7 @@ def run(args, config, defaults, config_file):
                             plot_dict,
                             "AlphaFold3",
                             plddt,
+                            pae,
                             score_file,
                             args.output_dir,
                         )
@@ -268,6 +288,7 @@ def run(args, config, defaults, config_file):
                         model.check_clashes()
                         score_file = bo.output[seed][idx]["json"]
                         plddt = model.residue_plddts
+                        pae = bo.output[seed][idx]["af3_pae"]
                         if len(indicies) > 0:
                             plddt = insert_none_by_minus_one(
                                 indicies[index_counter], plddt
@@ -278,6 +299,7 @@ def run(args, config, defaults, config_file):
                             plot_dict,
                             "Boltz",
                             plddt,
+                            pae,
                             score_file,
                             args.output_dir
                         )
@@ -294,6 +316,7 @@ def run(args, config, defaults, config_file):
                             model.check_clashes()
                             score_file = co.output[seed][idx]["scores"]
                             plddt = model.residue_plddts
+                            pae = co.output[seed][idx]["af3_pae"]
                             if len(indicies) > 0:
                                 plddt = insert_none_by_minus_one(
                                     indicies[index_counter], plddt
@@ -304,13 +327,45 @@ def run(args, config, defaults, config_file):
                                 plot_dict,
                                 "Chai-1",
                                 plddt,
+                                pae,
                                 score_file,
                                 args.output_dir,
                             )
                             chai_models["models"].append(model_data)
 
+        protenix_models = {"models": []}
+        if args.protenix:
+            if protenix_success:
+                programs_run.append("Protenix")
+                for seed in po.output.keys():
+                    for idx in po.output[seed].keys():
+                        if idx >= 0:
+                            model = po.output[seed][idx]["cif"]
+                            model.check_clashes()
+                            score_file = po.output[seed][idx]["scores"]
+                            plddt = model.residue_plddts
+                            pae = po.output[seed][idx]["af3_pae"]
+                            if len(indicies) > 0:
+                                plddt = insert_none_by_minus_one(
+                                    indicies[index_counter], plddt
+                                )
+                            index_counter += 1
+                            model_data = get_model_data(
+                                model,
+                                plot_dict,
+                                "Protenix",
+                                plddt,
+                                pae,
+                                score_file,
+                                args.output_dir,
+                            )
+                            protenix_models["models"].append(model_data)
+
         combined_models = (
-            alphafold_models["models"] + boltz_models["models"] + chai_models["models"]
+            alphafold_models["models"] +
+            boltz_models["models"] +
+            chai_models["models"] +
+            protenix_models["models"]
         )
 
         # Make the output directory for the models
@@ -324,6 +379,8 @@ def run(args, config, defaults, config_file):
                 output_name = "boltz_model_" + model["model_id"][-1] + ".cif"
             elif model["model_source"] == "Chai-1":
                 output_name = "chai_model_" + model["model_id"][-1] + ".cif"
+            elif model["model_source"] == "Protenix":
+                output_name = "protenix_model_" + model["model_id"][-1] + ".cif"
             shutil.copy(
                 cif_file,
                 args.output_dir.joinpath("output_models").joinpath(output_name),
@@ -414,11 +471,13 @@ def run(args, config, defaults, config_file):
 
 def main():
     """
-    Run AlphaFold3 / Boltz / Chai-1
+    Run AlphaFold3 / Boltz / Chai-1 / Protenix
     """
     import argparse
 
-    parser = argparse.ArgumentParser(description="Run AlphaFold3 / Boltz / Chai-1")
+    parser = argparse.ArgumentParser(
+        description="Run AlphaFold3 / Boltz / Chai-1 / Protenix"
+    )
 
     defaults = {}
     config_file = Path(__file__).parent.joinpath("data", "config.ini")
@@ -427,11 +486,13 @@ def main():
     if config_file.exists():
         config.read(str(config_file))
         defaults.update(dict(config.items("Databases")))
+        defaults.update(dict(config.items("Sif_paths")))
 
     parser = main_argpase_util(parser)
     parser = alphafold_argparse_util(parser)
     parser = boltz_argparse_util(parser)
     parser = chai_argparse_util(parser)
+    parser = protenix_argparse_util(parser)
     parser = mmseqs2_argparse_util(parser)
     parser = custom_template_argpase_util(parser)
     parser = prediction_argparse_util(parser)

abcfold/alphafold3/run_alphafold3.py

@@ -4,6 +4,8 @@
 from pathlib import Path
 from typing import Union
 
+from abcfold.alphafold3.check_install import check_af3_install
+
 logger = logging.getLogger("logger")
 
 
@@ -40,6 +42,9 @@ def run_alphafold3(
 
     input_json = Path(input_json)
     output_dir = Path(output_dir)
+
+    check_af3_install(interactive=False, sif_path=sif_path)
+
     cmd = generate_af3_cmd(
         input_json=input_json,
         output_dir=output_dir,