from . import util
import os
import numpy as np
import random
import sys
from alphafold.common import protein
from alphafold.model import data
from alphafold.model import config
from alphafold.model import model
from typing import Any, List, Mapping, NoReturn
from absl import logging
def set_config(
use_templates: bool,
max_msa_clusters: int,
max_extra_msa: int,
max_recycles: int,
model_id: int,
n_struct_module_repeats: int,
n_features_in: int,
monomer: bool = True,
model_params: int = 0,
) -> model.RunModel:
r"""Generated Runner object for AlphaFold
Parameters
----------
use_templates : Whether templates are used
max_msa_cluster : How many sequences to use in MSA
max_extra_msa : How many extra sequences to include for summary stats
max_recycles : Number of recycling iterations
model_id : Which AF2 model to use
n_struct_module_repeats : Number of passes through structure module
n_features_in : Unclear
monomer : Predicting as a monomer (set to False if using AlphaFold-multimer)
model_params : Which AF2 model config to use
Returns
----------
AlphaFold RunModel object
"""
if model_id not in range(1, 6):
logging.warning("model_id must be between 1 and 5!")
if use_templates:
model_id = random.randint(1, 2)
else:
model_id = random.randint(1, 5)
# Match model_params to model_id
# Sometimes we don't want to do this, for example,
# to reproduce output from ColabFold (which only uses models 1 and 3)
name = f"model_{ model_params }_ptm"
if not monomer:
name = f"model_{ model_params }_multimer"
cfg = config.model_config(name)
#### Provide config settings
#### MSAs
cfg.data.eval.num_ensemble = 1
if max_msa_clusters > 0:
cfg.data.eval.max_msa_clusters = min(n_features_in, max_msa_clusters)
if max_extra_msa > 0:
cfg.data.common.max_extra_msa = max(
1, min(n_features_in - max_msa_clusters, max_extra_msa)
)
#### Recycle and number of iterations
if monomer:
cfg.data.common.num_recycle = max_recycles
cfg.model.num_recycle = max_recycles
cfg.model.heads.structure_module.num_layer = n_struct_module_repeats
#### Templates
t = use_templates # for brevity
cfg.data.common.use_templates = use_templates
cfg.model.embeddings_and_evoformer.template.embed_torsion_angles = t
cfg.model.embeddings_and_evoformer.template.enabled = t
cfg.data.common.reduce_msa_clusters_by_max_templates = t
cfg.data.eval.subsample_templates = t
p = data.get_model_haiku_params(model_name=name, data_dir=".")
logging.debug("Prediction parameters:")
logging.debug("\tModel ID: {}".format(model_id))
logging.debug("\tUsing templates: {}".format(t))
logging.debug(
"\tMaximum MSA clusters: {}".format(cfg.data.eval.max_msa_clusters)
)
logging.debug(
"\tMaximum extra MSA clusters: {}".format(
cfg.data.common.max_extra_msa
)
)
logging.debug(
"\tNumber recycling iterations: {}".format(cfg.model.num_recycle)
)
logging.debug(
"\tNumber of structure module repeats: {}".format(
cfg.model.heads.structure_module.num_layer
)
)
return model.RunModel(cfg, p)
def run_one_job(
runner: model.RunModel, features_in: dict, random_seed: int, outname: str
) -> Mapping[str, Any]:
r"""Runs one AF2 job with input parameters
Parameters
----------
runner : AlphaFold2 job runner
features_in : Input features, including MSA and templates
random_seed : Random seed
outname : Name of PDB file to write
Returns
----------
None
"""
# Do one last bit of processing
features = runner.process_features(features_in, random_seed=random_seed)
# Generate the model
result = runner.predict(features, random_seed)
pred = protein.from_prediction(features, result)
# Write to file
to_pdb(outname, pred, result["plddt"], features_in["residue_index"])
return result
def predict_structure_from_templates(
seq: str,
outname: str,
a3m_lines: str,
template_path: str,
model_id: int = -1,
model_params: int = -1,
random_seed: int = -1,
max_msa_clusters: int = -1,
max_extra_msa: int = -1,
max_recycles: int = 3,
n_struct_module_repeats: int = 8,
) -> NoReturn:
r"""Predicts the structure.
Parameters
----------
seq : Sequence
outname : Name of output PDB
a3m_lines : String of entire alignment
template_paths : Where to locate templates
model_id : Which AF2 model to run (must be 1 or 2 for templates)
model_params : Which parameters to provide to AF2 model
random_seed : Random seed
max_msa_clusters : Number of sequences to use
max_extra_msa : Number of extra seqs for summary stats
max_recycles : Number of iterations through AF2
n_struct_module_repeats : Number of passes through structural refinement
move_prefix : Prefix for temporary files (deleted after fxn completion)
Returns
----------
None
"""
if random_seed == -1:
random_seed = random.randrange(sys.maxsize)
if model_id not in (1, 2):
model_id = random.randint(1, 2)
if model_params not in (1, 2):
model_params = random.randint(1, 2)
# Assemble the dictionary of input features
features_in = util.setup_features(
seq, a3m_lines, util.mk_template(seq, a3m_lines, template_path).features
)
# Run the models
model_runner = set_config(
True,
max_msa_clusters,
max_extra_msa,
max_recycles,
model_id,
n_struct_module_repeats,
len(features_in["msa"]),
model_params=model_params,
)
result = run_one_job(model_runner, features_in, random_seed, outname)
del model_runner
return result
def predict_structure_no_templates(
seq: str,
outname: str,
a3m_lines: str,
model_id: int = -1,
model_params: int = -1,
random_seed: int = -1,
max_msa_clusters: int = -1,
max_extra_msa: int = -1,
max_recycles: int = 3,
n_struct_module_repeats: int = 8,
) -> NoReturn:
r"""Predicts the structure.
Parameters
----------
seq : Sequence
outname : Name of output PDB
a3m_lines : String of entire alignment
model_id : Which AF2 model to run (must be 1 or 2 for templates)
random_seed : Random seed
max_msa_clusters : Number of sequences to use
max_extra_msa : Number of extra seqs for summary stats
max_recycles : Number of iterations through AF2
n_struct_module_repeats : Number of passes through structural refinement
Returns
----------
None
"""
# Set AF2 model details
if model_id not in range(1, 6):
model_id = random.randint(1, 5)
if model_params not in range(1, 6):
model_params = model_id
if random_seed == -1:
random_seed = random.randrange(sys.maxsize)
features_in = util.setup_features(seq, a3m_lines, util.mk_mock_template(seq))
model_runner = set_config(
False,
max_msa_clusters,
max_extra_msa,
max_recycles,
model_id,
n_struct_module_repeats,
len(features_in["msa"]),
model_params=model_params,
)
result = run_one_job(model_runner, features_in, random_seed, outname)
del model_runner
return result
def to_pdb(
outname, pred, plddts, res_idx # type unknown but check? # type unknown but check?
) -> NoReturn:
r"""Writes unrelaxed PDB to file
Parameters
----------
outname : Name of output PDB
pred : Prediction to write to PDB
plddts : Predicted errors
res_idx : Residues to print (default=all)
Returns
----------
None
"""
with open(outname, "w") as outfile:
outfile.write(protein.to_pdb(pred))
with open(f"b_{ outname }", "w") as outfile:
for line in open(outname, "r").readlines():
if line[0:6] == "ATOM ":
seq_id = int(line[22:26].strip()) - 1
seq_id = np.where(res_idx == seq_id)[0][0]
outfile.write(
"{}A{}{:6.2f}{}".format(
line[:21], line[22:60], plddts[seq_id], line[66:]
)
)
os.rename(f"b_{ outname }", outname)
