| import dataclasses
|
| import json
|
| from enum import Enum
|
| from typing import List, Dict, Tuple, Optional
|
|
|
|
|
| INTERFACE_MIN_ATOM_DIST = 8.0
|
|
|
|
|
| SubunitName = str
|
| ChainedSubunitName = str
|
| ChainName = str
|
| PdbPath = str
|
|
|
|
|
| @dataclasses.dataclass
|
| class SubunitInfo:
|
| name: SubunitName
|
| chain_names: List[str]
|
| start_res: int
|
| sequence: str
|
|
|
| def get_unstructured_res_ids(self) -> List[int]:
|
| return [self.start_res + i for i, resname in enumerate(self.sequence) if resname == "X"]
|
|
|
| def get_end_res(self) -> int:
|
| return self.start_res + len(self.sequence) - 1
|
|
|
| def get_active_res_ids(self) -> List[int]:
|
| return [i for i in range(self.start_res, self.get_end_res() + 1) if i not in self.get_unstructured_res_ids()]
|
|
|
| def get_relative_active_res_ids(self) -> List[int]:
|
| return [i - self.start_res for i in range(self.start_res, self.get_end_res() + 1)
|
| if i not in self.get_unstructured_res_ids()]
|
|
|
| def get_chained_names(self) -> List[str]:
|
| return [f"{self.name}_{chain_name}" for chain_name in self.chain_names]
|
|
|
| def to_dict(self):
|
| return dataclasses.asdict(self)
|
|
|
| @classmethod
|
| def from_dict(cls, d):
|
| return cls(**d)
|
|
|
|
|
| SubunitsInfo = Dict[SubunitName, SubunitInfo]
|
|
|
|
|
| def save_subunits_info(subunits_info: SubunitsInfo, output_path: str):
|
| json_data = {domain_name: domain_info.to_dict() for domain_name, domain_info in subunits_info.items()}
|
| json.dump(json_data, open(output_path, "w"), indent=2)
|
|
|
|
|
| def read_subunits_info(output_path: str) -> SubunitsInfo:
|
| json_data = json.load(open(output_path))
|
|
|
| return {domain_name: SubunitInfo.from_dict({k: v for k, v in domain_info.items() if k != "end_res"})
|
| for domain_name, domain_info in json_data.items()}
|
|
|
|
|
| @dataclasses.dataclass
|
| class AlphaFoldJobInfo:
|
| subunit_names: List[SubunitName]
|
| merged_subunits: List[bool]
|
| sequences: List[str]
|
|
|
| def get_jobname(self):
|
| return "_".join(self.subunit_names)
|
|
|
| def get_as_fasta(self):
|
| return f">{self.get_jobname()}\n" + ":".join(self.sequences) + "\n"
|
|
|
| def __hash__(self):
|
| return hash(tuple([*self.subunit_names, *self.merged_subunits]))
|
|
|
|
|
| class RunAlphaFoldResult(Enum):
|
| SUCCESS = 0
|
| NOT_STARTED = 1
|
| RUNNING = 2
|
| FAILED = 3
|
| SHOULD_RERUN = 4
|
|
|
|
|
| @dataclasses.dataclass
|
| class AFSubunitScores:
|
| plddt_avg: float
|
| plddt_percentile: List[float]
|
| plddt_interface_avg: float
|
| plddt_interface_percentile: List[float]
|
|
|
| self_pae_avg: float
|
| self_pae_percentile: List[float]
|
|
|
|
|
| @dataclasses.dataclass
|
| class AFInteractionScores:
|
| pae_avg: float
|
| pae_percentile: List[float]
|
| pae_joined_interface_avg: float
|
| pae_joined_interface_percentile: List[float]
|
|
|
| interface1_size: int
|
| interface2_size: int
|
|
|
|
|
| @dataclasses.dataclass
|
| class SubunitPdbInfo:
|
| chain_id: ChainName
|
| chain_residue_id: int
|
| pdb_residue_id: int
|
| length: int
|
|
|
|
|
| @dataclasses.dataclass
|
| class AFResultScoredPair:
|
| subunits_names: Tuple[SubunitName, SubunitName]
|
| pdb_path: str
|
|
|
| subunit1_pdb_info: SubunitPdbInfo
|
| subunit2_pdb_info: SubunitPdbInfo
|
|
|
|
|
|
|
| subunit1_scores: Optional[AFSubunitScores] = None
|
| subunit2_scores: Optional[AFSubunitScores] = None
|
| interaction_scores: Optional[AFInteractionScores] = None
|
|
|
