File size: 2,933 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 | import dataclasses
import json
from enum import Enum
from typing import List, Dict, Tuple, Optional
# Constants
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 # inclusive, should be greater than 0
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
|