| from collections import defaultdict
|
| from typing import List, Optional, Dict
|
|
|
| from automatic_pipeline.configurable import MAX_AF_SIZE
|
| from automatic_pipeline.libs.utils_classes import SubunitsInfo, AlphaFoldJobInfo, SubunitInfo, AFResultScoredPair, \
|
| SubunitName
|
|
|
|
|
| def get_residue_diff(subunit1: SubunitInfo, subunit2: SubunitInfo) -> int:
|
| return (subunit2.start_res - 1) - subunit1.get_end_res()
|
|
|
|
|
| def get_merged_info(subunits_info: SubunitsInfo, subunit_names: List[str]) -> List[bool]:
|
| ret_merged = []
|
| residues_less_than_max = MAX_AF_SIZE - sum([len(subunits_info[s].sequence) for s in subunit_names])
|
|
|
| for i in range(len(subunit_names) - 1):
|
| subunit1 = subunits_info[subunit_names[i]]
|
| subunit2 = subunits_info[subunit_names[i + 1]]
|
|
|
| residue_diff = get_residue_diff(subunit1, subunit2)
|
| if subunit1.chain_names[0] == subunit2.chain_names[0] \
|
| and 0 <= residue_diff <= min(residues_less_than_max, 100) and subunit1.name != subunit2.name:
|
| residues_less_than_max -= residue_diff
|
| ret_merged.append(True)
|
| else:
|
| ret_merged.append(False)
|
| return ret_merged
|
|
|
|
|
| def get_alphafold_job(subunits_info: SubunitsInfo, subunit_names: List[str]) -> Optional[AlphaFoldJobInfo]:
|
| assert tuple(subunit_names) == tuple(sorted(subunit_names)), f"subunit names must be sorted {subunit_names}"
|
| sequences_length = sum([len(subunits_info[subunit_name].sequence) for subunit_name in subunit_names])
|
| if sequences_length > MAX_AF_SIZE:
|
| print(f"Skipping because too long", subunit_names)
|
| return None
|
|
|
| merged_subunits = get_merged_info(subunits_info, subunit_names)
|
|
|
| sequences = []
|
| next_sequence = subunits_info[subunit_names[0]].sequence
|
|
|
| for i in range(1, len(subunit_names)):
|
| if merged_subunits[i - 1]:
|
| residue_diff = get_residue_diff(subunits_info[subunit_names[i - 1]], subunits_info[subunit_names[i]])
|
| next_sequence += "X" * residue_diff
|
| next_sequence += subunits_info[subunit_names[i]].sequence
|
| else:
|
| sequences.append(next_sequence)
|
| next_sequence = subunits_info[subunit_names[i]].sequence
|
|
|
| if next_sequence:
|
| sequences.append(next_sequence)
|
|
|
| return AlphaFoldJobInfo(subunit_names=subunit_names,
|
| merged_subunits=get_merged_info(subunits_info, subunit_names),
|
| sequences=sequences)
|
|
|
|
|
| def get_af_jobs_for_pairs(subunits_info: SubunitsInfo) -> List[AlphaFoldJobInfo]:
|
| returned_jobs = []
|
| all_subunits_names = sorted(list(subunits_info.keys()))
|
| for i in range(len(all_subunits_names)):
|
| start_from = i if len(subunits_info[all_subunits_names[i]].chain_names) > 1 else i + 1
|
| for j in range(start_from, len(all_subunits_names)):
|
| af_job = get_alphafold_job(subunits_info, [all_subunits_names[i], all_subunits_names[j]])
|
| if af_job is not None:
|
| returned_jobs.append(af_job)
|
|
|
| return returned_jobs
|
|
|
|
|
| def _score_result_transform(result: AFResultScoredPair) -> float:
|
| full_pae_avg: float = (result.interaction_scores.pae_avg + result.subunit1_scores.self_pae_avg
|
| + result.subunit2_scores.self_pae_avg) / 3
|
|
|
|
|
| pae_based_score = max([1, 100 - (1/4) * (full_pae_avg ** 2)])
|
| iplddt_score = (result.subunit1_scores.plddt_interface_avg + result.subunit2_scores.plddt_interface_avg) / 2
|
| return pae_based_score + iplddt_score/100
|
|
|
|
|
| def get_job_length(subunit_names, subunits_info):
|
| return sum([len(subunits_info[subunit_name].sequence) for subunit_name in subunit_names])
|
|
|
|
|
| def get_af_jobs_for_groups(subunits_info: SubunitsInfo, parsed_pairs: List[AFResultScoredPair])\
|
| -> List[AlphaFoldJobInfo]:
|
| groups_jobs = set()
|
|
|
| best_for_subunit: Dict[SubunitName, Dict[SubunitName, float]] = defaultdict(dict)
|
| for pair in parsed_pairs:
|
| subunit1, subunit2 = pair.subunits_names
|
| score = _score_result_transform(pair)
|
|
|
| best_for_subunit[subunit1][subunit2] = max(best_for_subunit[subunit1].get(subunit2, 0), score)
|
| best_for_subunit[subunit2][subunit1] = max(best_for_subunit[subunit2].get(subunit1, 0), score)
|
| best_for_subunit = dict(best_for_subunit)
|
|
|
| for subunit_name, subunit_info in subunits_info.items():
|
| sorted_best_for_subunit = sorted(best_for_subunit[subunit_name].keys(),
|
| key=lambda x: best_for_subunit[subunit_name][x], reverse=True)
|
|
|
| print("best for", subunit_name, best_for_subunit[subunit_name], sorted_best_for_subunit)
|
|
|
| job = [subunit_name, sorted_best_for_subunit[0]]
|
| for i in range(3, 6):
|
| for subunit_to_add in sorted_best_for_subunit:
|
| if job.count(subunit_to_add) < len(subunits_info[subunit_to_add].chain_names) \
|
| and get_job_length(job + [subunit_to_add], subunits_info) <= MAX_AF_SIZE:
|
| job.append(subunit_to_add)
|
| break
|
| if len(job) != i:
|
| break
|
| print("not adding yet, with repetitions", job)
|
| groups_jobs.add(tuple(sorted(job)))
|
|
|
|
|
|
|
|
|
|
|
| job = [subunit_name, sorted_best_for_subunit[0]]
|
| for subunit_to_add in sorted_best_for_subunit[1:]:
|
| job.append(subunit_to_add)
|
| if get_job_length(job, subunits_info) > MAX_AF_SIZE:
|
| break
|
| if len(job) > 5:
|
| break
|
| print("not adding yet, without repetitions", job)
|
| groups_jobs.add(tuple(sorted(job)))
|
|
|
|
|
|
|
|
|
|
|
| return [get_alphafold_job(subunits_info, list(job)) for job in groups_jobs]
|
|
|
