File size: 8,098 Bytes
44459bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 |
"""AF2 folding submission command."""
from pathlib import Path
from typing import List, Optional
import typer
from folding_studio_data_models import (
FeatureMode,
)
from folding_studio_data_models.request.folding import FoldingModel
from typing_extensions import Annotated
from folding_studio.api_call.predict import (
batch_prediction,
batch_prediction_from_file,
simple_prediction,
)
from folding_studio.commands.predict.utils import (
print_instructions_batch,
print_instructions_simple,
validate_model_subset,
validate_source_path,
)
from folding_studio.config import FOLDING_API_KEY
from folding_studio.console import console
from folding_studio.utils.data_model import (
BatchInputFile,
PredictRequestCustomFiles,
PredictRequestParams,
)
from folding_studio.utils.input_validation import (
extract_and_validate_custom_msas,
extract_and_validate_custom_templates,
validate_initial_guess,
)
def af2( # pylint: disable=dangerous-default-value, too-many-arguments, too-many-locals
source: Annotated[
Path,
typer.Argument(
help=(
"Path to the data source. Either a fasta file, a directory of fasta files "
"or a csv/json file describing a batch prediction request."
),
callback=validate_source_path,
exists=True,
),
],
project_code: Annotated[
str,
typer.Option(
help=(
"Project code. If unknown, contact your PM or the Folding Studio team."
),
exists=True,
envvar="FOLDING_PROJECT_CODE",
),
],
cache: Annotated[
bool,
typer.Option(help="Use cached experiment results if any."),
] = True,
template_mode: Annotated[
FeatureMode,
typer.Option(help="Mode of the template features generation."),
] = FeatureMode.SEARCH,
custom_template: Annotated[
List[Path],
typer.Option(
help=(
"Path to a custom template or a directory of custom templates. "
"To pass multiple inputs, simply repeat the flag "
"(e.g. `--custom_template template_1.cif --custom_template template_2.cif`)."
),
callback=extract_and_validate_custom_templates,
exists=True,
),
] = [],
custom_template_id: Annotated[
List[str],
typer.Option(
help=(
"ID of a custom template. "
"To pass multiple inputs, simply repeat the flag "
"(e.g. `--custom_template_id template_ID_1 --custom_template_id template_ID_2`)."
)
),
] = [],
initial_guess_file: Annotated[
Path | None,
typer.Option(
help=("Path to an initial guess file."),
callback=validate_initial_guess,
exists=True,
),
] = None,
templates_masks_file: Annotated[
Path | None,
typer.Option(
help=("Path to a templates masks file."),
exists=True,
),
] = None,
msa_mode: Annotated[
FeatureMode,
typer.Option(help="Mode of the MSA features generation."),
] = FeatureMode.SEARCH,
custom_msa: Annotated[
List[Path],
typer.Option(
help=(
"Path to a custom msa or a directory of custom msas. "
"To pass multiple inputs, simply repeat the flag "
"(e.g. `--custom_msa msa_1.sto --custom_msa msa_2.sto`)."
),
callback=extract_and_validate_custom_msas,
exists=True,
),
] = [],
max_msa_clusters: Annotated[
int,
typer.Option(help="Max number of MSA clusters to search."),
] = -1,
max_extra_msa: Annotated[
int,
typer.Option(
help="Max extra non-clustered MSA representation to use as source."
),
] = -1,
gap_trick: Annotated[
bool,
typer.Option(
help="Activate gap trick, allowing to model complexes with monomer models."
),
] = False,
num_recycle: Annotated[
int,
typer.Option(
help="Number of refinement iterations of the predicted structures."
),
] = 3,
model_subset: Annotated[
list[int],
typer.Option(
help="Subset of AF2 model ids to use, between 1 and 5 included.",
callback=validate_model_subset,
),
] = [],
random_seed: Annotated[
int,
typer.Option(
help=(
"Random seed used during the MSA sampling. "
"Different random seed values will introduce variations in the predictions."
)
),
] = 0,
num_seed: Annotated[
Optional[int],
typer.Option(
help="Number of random seeds to use. Creates a batch prediction.", min=2
),
] = None,
metadata_file: Annotated[
Optional[Path],
typer.Option(
help=(
"Path to the file where the job metadata returned by the server are written."
),
),
] = None,
):
"""Asynchronous AF2 folding submission.
Read more at https://int-bio-foldingstudio-gcp.nw.r.appspot.com/how-to-guides/af2_openfold/single_af2_job/.
If the source is a CSV or JSON file describing a batch prediction request, all the other
options will be overlooked.
"""
if FOLDING_API_KEY:
console.print(":key: Using detected API key for authentication.")
else:
console.print(":yellow_circle: Using JWT for authentication.")
is_batch = source.is_dir() or source.suffix in BatchInputFile.__members__.values()
is_multi_seed = num_seed is not None
is_batch = is_batch or is_multi_seed
params = PredictRequestParams(
ignore_cache=not cache,
template_mode=template_mode,
custom_template_ids=custom_template_id,
msa_mode=msa_mode,
max_msa_clusters=max_msa_clusters,
max_extra_msa=max_extra_msa,
gap_trick=gap_trick,
num_recycle=num_recycle,
random_seed=random_seed,
model_subset=model_subset,
)
custom_files = PredictRequestCustomFiles(
templates=custom_template,
msas=custom_msa,
initial_guess_files=[initial_guess_file] if initial_guess_file else None,
templates_masks_files=[templates_masks_file] if templates_masks_file else None,
)
if is_batch:
if is_multi_seed:
response = batch_prediction(
files=[source],
folding_model=FoldingModel.AF2,
params=params,
custom_files=custom_files,
num_seed=num_seed,
project_code=project_code,
)
elif source.is_file():
console.print(
f"Submitting batch jobs configuration file [bold]{source}[/bold]"
)
console.print(
"Input options are [bold yellow]ignored[/bold yellow] in favor of the configuration file content."
)
response = batch_prediction_from_file(
file=source,
project_code=project_code,
)
elif source.is_dir():
response = batch_prediction(
files=list(f for f in source.iterdir() if f.is_file()),
folding_model=FoldingModel.AF2,
params=params,
custom_files=custom_files,
num_seed=num_seed,
project_code=project_code,
)
print_instructions_batch(response_json=response, metadata_file=metadata_file)
else:
response = simple_prediction(
file=source,
folding_model=FoldingModel.AF2,
params=params,
custom_files=custom_files,
project_code=project_code,
)
print_instructions_simple(response_json=response, metadata_file=metadata_file)
|