File size: 2,824 Bytes
d149fb1 | 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 | import argparse
import os
import pickle
from onescience.flax_models.alphafold.data import pipeline, pipeline_multimer, templates
from onescience.flax_models.alphafold.data.tools import hmmsearch, hhsearch
from utils import add_data_args
def main(args):
if (args.multimer):
template_searcher = hmmsearch.Hmmsearch(
binary_path=args.hmmsearch_binary_path,
hmmbuild_binary_path=args.hmmbuild_binary_path,
database_path=args.pdb_seqres_database_path,
)
template_featurizer = templates.HmmsearchHitFeaturizer(
mmcif_dir=args.template_mmcif_dir,
max_template_date=args.max_template_date,
max_hits=20,
kalign_binary_path=args.kalign_binary_path,
release_dates_path=args.release_dates_path,
obsolete_pdbs_path=args.obsolete_pdbs_path
)
else:
template_searcher = hhsearch.HHSearch(
binary_path=args.hhsearch_binary_path,
databases=[args.pdb70_database_path],
)
template_featurizer = templates.HhsearchHitFeaturizer(
mmcif_dir=args.template_mmcif_dir,
max_template_date=args.max_template_date,
max_hits=20,
kalign_binary_path=args.kalign_binary_path,
release_dates_path=None,
obsolete_pdbs_path=args.obsolete_pdbs_path
)
data_pipeline = pipeline.DataPipeline(
jackhmmer_binary_path=args.jackhmmer_binary_path,
hhblits_binary_path=args.hhblits_binary_path,
uniref90_database_path=args.uniref90_database_path,
mgnify_database_path=args.mgnify_database_path,
bfd_database_path=args.bfd_database_path,
uniref30_database_path=args.uniref30_database_path,
small_bfd_database_path=None,
template_featurizer=template_featurizer,
template_searcher=template_searcher,
use_small_bfd=False,
)
if (args.multimer):
data_pipeline = pipeline_multimer.DataPipeline(
monomer_data_pipeline=data_pipeline,
jackhmmer_binary_path=args.jackhmmer_binary_path,
uniprot_database_path=args.uniprot_database_path)
feature_dict = data_pipeline.process(
input_fasta_path=args.fasta_path,
msa_output_dir=args.output_dir,
)
with open(os.path.join(args.output_dir, "feature_dict.pickle"), "wb") as fp:
pickle.dump(feature_dict, fp, protocol=pickle.HIGHEST_PROTOCOL)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("fasta_path", type=str)
parser.add_argument("mmcif_dir", type=str)
parser.add_argument("output_dir", type=str)
parser.add_argument("--multimer", action='store_true')
add_data_args(parser)
args = parser.parse_args()
main(args)
|