Add model_preset argument/ option/ parameter

run_af2/bin/submit-af2

@@ -28,15 +28,17 @@
 if test $# -lt 2; then
     echo "usage: submit-af2 [--use-gpu]"
     echo "                  [--db-preset {reduced_dbs,full_dbs}]"
+    echo "                  [model-preset {monomer,monomer_casp14,monomer_ptm,multimer}]"
     echo "                  [--max-template-date YYYY-MM-DD]"
     echo "                  <OUTPUT DIR> <FASTA FILE> [<FASTA FILE> ...]"
     exit 1
 fi
 
-# There are only two non-positional arguments, simply loop over candidates
+# There are only two non-positional arguments, simply loop over candidates.
+# This should be turned into switch..case if argument extension goes on.
 SLURM_AF_PARAMS="--exclude=lii[02-28]"
 AF_PIPELINE_PARAM=""
-for i in 1 2 3; do
+for i in 1 2 3 4; do
     if test x"${1}" = x"--use-gpu"; then
         shift
         SLURM_AF_PARAMS="--partition=a100,rtx8000 --gres=gpu:1"
@@ -49,6 +51,10 @@ for i in 1 2 3; do
                   shift
                   AF_PIPELINE_PARAM="${AF_PIPELINE_PARAM} --max-template-date ${1}"
                   shift
+              else if test x"${1}" = x"--model-preset"; then
+                       shift
+                       AF_PIPELINE_PARAM="${AF_PIPELINE_PARAM} --model-preset ${1}"
+                   fi
               fi
          fi
     fi

run_af2/src/run_af2/_shared.py

@@ -28,6 +28,8 @@ functionality to be external, we should create a proper module for that.
 import os
 import sys
 
+MODEL_PRESET_CHOICES = ["monomer", "monomer_casp14", "monomer_ptm", "multimer"]
+
 
 def get_version():
     """Load the package version (import only upon request).
@@ -96,12 +98,18 @@ def parse_af2_arguments(parser):
         "-p",
         "--db-preset",
         choices=db_preset_choices.keys(),
-        help="Choose model configuration - no ensembling and smaller genetic "
-        + "databases (reduced_dbs), no ensembling and full genetic databases "
-        + "configuration (full_dbs), the casp14 preset is not available right "
-        + "now (ask sciCORE to install the data).",
+        help="Choose databases - smaller genetic databases (reduced_dbs) or "
+        + "full genetic databases (full_dbs).",
         default="full_dbs",
     )
+    af2_group.add_argument(
+        "-m",
+        "--model-preset",
+        choices=MODEL_PRESET_CHOICES,
+        help="Choose the model to be run, follows the --model_preset of the "
+        + "AF2 pipeline.",
+        default="monomer",
+    )
     af2_group.add_argument(
         "--version",
         action="version",

run_af2/src/run_af2/run_singularity.py

@@ -115,6 +115,7 @@ def _assemble_singularity_call(  # pylint: disable=too-many-arguments
     binds,
     use_gpu,
     db_preset,
+    model_preset,
     data_paths,
 ):
     """Assemble the command to run AF2 from the Singularity image."""
@@ -146,6 +147,7 @@ def _assemble_singularity_call(  # pylint: disable=too-many-arguments
             f"--output_dir={snglrty_out}",
             f"--max_template_date={max_template_date}",
             f"--db_preset={db_preset}",
+            f"--model_preset={model_preset}",
             "--logtostderr",
         ]
     )
@@ -186,6 +188,7 @@ def run_af2_singularity_image(  # pylint: disable=too-many-arguments
     max_template_date=None,
     use_gpu=False,
     db_preset="full_dbs",
+    model_preset="monomer",
     af2_image_file=None,
     # "/export/soft/singularity-containers/alphafold",
     af2_image_dir="/scicore/home/schwede/GROUP/alphafold_data/",
@@ -230,11 +233,13 @@ def run_af2_singularity_image(  # pylint: disable=too-many-arguments
     :type max_template_date: :class:`str`
     :param use_gpu: Run the AF2 pipeline using GPUs or not.
     :type use_gpu: :class:`bool`
-    :param db_preset: Model configuration - no ensembling and smaller genetic
-                      databases (reduced_dbs), no ensembling and full genetic
-                      databases configuration (full_dbs). Corresponds to the
-                      preset parameter in AF2.
+    :param db_preset: Database configuration - smaller genetic databases
+                      (reduced_dbs) and full genetic databases configuration
+                      (full_dbs). Corresponds to the db_preset parameter in AF2.
     :type db_preset: :class:`str`
+    :param model_preset: Model configuration - corresponds to the model_preset
+                         parameter in AF2.
+    :type model_preset: :class:`str`
     :param af2_image_file: Declare a Singularity image to run the AF2 pipeline
                            from. If None, an image from af2_image_dir will be
                            used.
@@ -296,6 +301,11 @@ def run_af2_singularity_image(  # pylint: disable=too-many-arguments
             f"db_preset {db_preset} not allowed, known values: "
             + "'full_dbs', 'reduced_dbs'"
         )
+    if model_preset not in _shared.MODEL_PRESET_CHOICES:
+        raise ValueError(
+            f"model_reset {model_preset} not allowed, known "
+            + f"values: {', '.join(_shared.MODEL_PRESET_CHOICES)}"
+        )
 
     # gather data paths into dictionary for mountpoint collecting
     data_paths = {
@@ -337,6 +347,7 @@ def run_af2_singularity_image(  # pylint: disable=too-many-arguments
         binds,
         use_gpu,
         db_preset,
+        model_preset,
         data_paths,
     )
 
@@ -399,6 +410,7 @@ def main():
             opts.max_template_date,
             opts.use_gpu,
             opts.db_preset,
+            opts.model_preset,
             af2_image_file=opts.singularity_image,
             **data_paths,
         )