diff --git a/run_af2/bin/submit-af2 b/run_af2/bin/submit-af2
index 9b19a76c3fc9a639f67707f941bc630405de3bbe..7d5b1d075d14991bf964d401dafeda682cf5bd0a 100755
--- a/run_af2/bin/submit-af2
+++ b/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
diff --git a/run_af2/src/run_af2/_shared.py b/run_af2/src/run_af2/_shared.py
index eee1d955964b6df6dfe5d3a846720e27a0b38ae4..5d7f22c75261a5320142ae8b6e69550fdc399616 100644
--- a/run_af2/src/run_af2/_shared.py
+++ b/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",
diff --git a/run_af2/src/run_af2/run_singularity.py b/run_af2/src/run_af2/run_singularity.py
index 165f788753dcfb798ec5b2528d10102bcbcb5574..5d87fa3b990016e6304fd932c54e442cb3aea9a5 100644
--- a/run_af2/src/run_af2/run_singularity.py
+++ b/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,
)