allowing for distogram outputs by alphafold with abcfold

.gitignore

@@ -13,7 +13,6 @@ build/
 
 .history/
 
-
 # pip packages
 dist/
 *.egg-info/

abcfold/abcfold.py

@@ -1,6 +1,7 @@
 import configparser
 import json
 import os
+import re
 import shutil
 import socketserver
 import sys
@@ -153,14 +154,16 @@ def run(args, config, defaults, config_file):
                 number_of_models=args.number_of_models,
                 num_recycles=args.num_recycles,
                 sif_path=args.sif_path,
+                save_distogram=args.save_distogram
             )
 
             if af3_success:
+                pattern = re.compile(f"^{name}.*", re.IGNORECASE)
                 af3_out_dir = list(
                     [
                         dir_
-                        for dir_ in args.output_dir.glob(f"*{name.lower()}*")
-                        if dir_.is_dir()
+                        for dir_ in args.output_dir.glob("*")
+                        if dir_.is_dir() and pattern.match(dir_.name)
                     ]
                 )[0]
                 ao = AlphafoldOutput(af3_out_dir, input_params, name)

abcfold/alphafold3/run_alphafold3.py

@@ -16,6 +16,7 @@ def run_alphafold3(
     interactive: bool = False,
     number_of_models: int = 5,
     num_recycles: int = 10,
+    save_distogram: bool = False
 ) -> bool:
     """
     Run Alphafold3 using the input JSON file
@@ -48,6 +49,7 @@ def run_alphafold3(
         interactive=interactive,
         number_of_models=number_of_models,
         num_recycles=num_recycles,
+        save_distogram=save_distogram
     )
 
     logger.info("Running Alphafold3")
@@ -77,6 +79,7 @@ def generate_af3_cmd(
     number_of_models: int = 10,
     num_recycles: int = 5,
     interactive: bool = False,
+    save_distogram: bool = False,
 ) -> str:
     """
     Generate the Alphafold3 command
@@ -110,7 +113,8 @@ def generate_af3_cmd(
         --model_dir=/root/models \
         --output_dir=/root/af_output \
         --num_diffusion_samples {number_of_models}\
-        --num_recycles {num_recycles}
+        --num_recycles {num_recycles}\
+        --save_distogram {str(save_distogram).lower()}
     """
 
     else:
@@ -127,5 +131,6 @@ def generate_af3_cmd(
         --model_dir=/root/models \
         --output_dir=/root/af_output \
         --num_diffusion_samples {number_of_models}\
-        --num_recycles {num_recycles}
+        --num_recycles {num_recycles}\
+        --save_distogram {str(save_distogram).lower()}
         """

abcfold/argparse_utils.py

@@ -174,6 +174,12 @@ def alphafold_argparse_util(parser):
 Alphafold3 to search for templates",
     )
 
+    parser.add_argument(
+        "--save_distogram",
+        action="store_true",
+        help="[optional] store AlphaFold3 distograms",
+    )
+
     return parser