| 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 = {subunit_name: subunit_info.to_dict() for subunit_name, subunit_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))
|
| subunits_info = {subunit_name: SubunitInfo.from_dict({k: v for k, v in subunit_info.items() if k != "end_res"})
|
| for subunit_name, subunit_info in json_data.items()}
|
| for subunit_name, subunit_info in subunits_info.items():
|
| assert subunit_name == subunit_info.name, f"Key name and subunit info name must match, mismatch:" \
|
| f" {subunit_name} and {subunit_info.name}"
|
| res_to_subunit_name = {}
|
| for subunit_name, subunit_info in subunits_info.items():
|
| for chain_name in subunit_info.chain_names:
|
| for res_id in subunit_info.get_active_res_ids():
|
| res_key = (chain_name, res_id)
|
| assert res_key not in res_to_subunit_name, \
|
| f"Residue {res_id} in chain {chain_name} is present in multiple subunits - " \
|
| f"{subunit_name} and {res_to_subunit_name[res_key]}. overlaps are not allowed in subunits."
|
| res_to_subunit_name[res_key] = subunit_name
|
| return subunits_info
|
|
|
|
|
| @dataclasses.dataclass
|
| class SubunitPdbInfo:
|
| chain_id: ChainName
|
| chain_residue_id: int
|
| pdb_residue_id: int
|
| length: int
|
|
|
