File size: 4,329 Bytes
8efb4bd | 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 | import os
import sys
from collections import Counter
from libs.get_alphafold_jobs import get_af_jobs_for_pairs, get_af_jobs_for_groups
from libs.parse_alphafold_jobs import score_af_results_as_pairs
from libs.run_alphafold_jobs import run_alphafold_on_jobs
from libs.run_assembly import run_assembly, get_assembly_results
from libs.utils_classes import RunAlphaFoldResult, read_subunits_info
THIS_SCRIPT_PATH = os.path.abspath(__file__)
BASE_PATH = os.path.dirname(THIS_SCRIPT_PATH)
SUBUNITS_INFO_PATH = os.path.join(BASE_PATH, "subunits_info")
MAIN_OUTPUT_PATH = os.path.join(BASE_PATH, "output")
def run_on_subunits_json(subunits_json_path: str, output_path: str):
subunits_info = read_subunits_info(subunits_json_path)
# run alphafold pairs
print("-- Pairs run AlphaFold")
pairs_af_output_path = os.path.join(output_path, "alphafold_pairs")
pairs_fasta_path = os.path.join(output_path, "fastas_pairs")
pairs_alphafold_jobs = get_af_jobs_for_pairs(subunits_info)
alphafold_results = run_alphafold_on_jobs(pairs_alphafold_jobs, pairs_fasta_path, pairs_af_output_path)
results_counter = Counter([i[0] for i in alphafold_results.values()])
if results_counter[RunAlphaFoldResult.RUNNING] > 0:
print(f"still running", results_counter)
return "af_pairs"
# parse pairs
print("-- Pairs Parse AlphaFold")
parsed_pairs_path = os.path.join(output_path, "parsed_pairs.pkl")
parsed_pairs = score_af_results_as_pairs({af_job: pdb_paths for af_job, (_, pdb_paths) in
alphafold_results.items()}, parsed_pairs_path, subunits_info)
# run alphafold groups
print("-- Groups run AlphaFold")
groups_af_output_path = os.path.join(output_path, "alphafold_groups")
groups_fasta_path = os.path.join(output_path, "fastas_groups")
groups_alphafold_jobs = get_af_jobs_for_groups(subunits_info, parsed_pairs)
print("Created groups", [i.subunit_names for i in groups_alphafold_jobs])
return "tmp_no_groups"
alphafold_results = run_alphafold_on_jobs(groups_alphafold_jobs, groups_fasta_path, groups_af_output_path)
results_counter = Counter([i[0] for i in alphafold_results.values()])
if results_counter[RunAlphaFoldResult.RUNNING] > 0:
print(f"still running", results_counter)
return "af_groups"
# parse groups
print("-- Groups Parse AlphaFold")
parsed_groups_path = os.path.join(output_path, "parsed_groups.pkl")
parsed_groups = score_af_results_as_pairs({af_job: pdb_paths for af_job, (_, pdb_paths) in
alphafold_results.items()}, parsed_groups_path, subunits_info)
# run Hierarchical Assembly of Predicted Subunits (HAPS)
print("-- Assembly")
assembly_path = os.path.join(output_path, "assembly")
assembly_result = run_assembly(parsed_pairs + parsed_groups, subunits_info, assembly_path)
if assembly_result == "running":
print(f"still running")
return "assembly"
elif assembly_result == "success":
print("done_" + str(get_assembly_results(assembly_path)))
return "done_" + str(get_assembly_results(assembly_path))
return "done_0"
def run_on_folders():
progression = {"done": [], "af_pairs": [], "af_groups": [], "assembly": []}
scores = []
for subunits_info_json_path in os.listdir(SUBUNITS_INFO_PATH):
jobname = subunits_info_json_path.split(".")[0]
print("--------", jobname, "--------")
stage = run_on_subunits_json(os.path.join(SUBUNITS_INFO_PATH, subunits_info_json_path),
os.path.join(MAIN_OUTPUT_PATH, jobname))
if stage.startswith("done_"):
scores.append((jobname, float(stage.split("_")[1])))
stage = "done"
progression[stage].append(jobname)
print({f"{k}_{len(v)}": v for k, v in progression.items()})
print(scores)
if __name__ == '__main__':
if len(sys.argv) == 2 and sys.argv[1] == "all":
run_on_folders()
elif len(sys.argv) != 3:
print("usage: <script> subunits_info output_path")
else:
run_on_subunits_json(os.path.abspath(sys.argv[1]), os.path.abspath(sys.argv[2]))
|