| import os
|
| import time
|
| from collections import defaultdict
|
| from typing import Dict, Tuple, List
|
|
|
| import Bio.PDB
|
| import Bio.SeqIO
|
| import numpy as np
|
| import scipy.spatial
|
| from scipy.spatial import KDTree
|
|
|
| INTERFACE_MIN_ATOM_DIST = 5
|
|
|
|
|
| def get_chain_to_seq(pdb_path: str) -> Dict[str, str]:
|
| chain_to_seq = {str(record.id): str(record.seq) for record in Bio.SeqIO.parse(pdb_path, 'pdb-seqres')}
|
| if len(chain_to_seq) > 0:
|
| return chain_to_seq
|
|
|
|
|
| def create_ident_chain_map(chain_to_seq: Dict[str, str]) -> Dict[str, str]:
|
| seq_to_chains = defaultdict(list)
|
| for chain_name, seq in chain_to_seq.items():
|
| seq_to_chains[seq].append(chain_name)
|
|
|
| ident_chain_map = {}
|
| for v in seq_to_chains.values():
|
| ident_name = sorted(v)[0]
|
| for chain_name in v:
|
| ident_chain_map[chain_name] = ident_name
|
|
|
| return ident_chain_map
|
|
|
|
|
| def get_ident_chain_map_from_complex(pdb_path: str) -> Dict[str, str]:
|
| return create_ident_chain_map(get_chain_to_seq(pdb_path))
|
|
|
|
|
| def get_contacts(pdb_path: str, ident_chains: Dict[str, str], sample_chain_map=None, offset: int = 0,
|
| bfactor_threshold=0):
|
| pdb_parser = Bio.PDB.PDBParser(QUIET=True)
|
| pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
|
| pdb_model = next(iter(pdb_struct))
|
|
|
| coords = []
|
| identifiers: List[Tuple[str, int]] = []
|
| for res in pdb_model.get_residues():
|
| if "CA" in res:
|
| if res["CA"].get_bfactor() < bfactor_threshold:
|
| continue
|
| for atom in res:
|
|
|
| if atom.element == "H":
|
| continue
|
| coords.append(atom.get_coord())
|
| if sample_chain_map is None:
|
|
|
| identifiers.append((res.parent.id, res.id[1] + offset))
|
| else:
|
|
|
| identifiers.append((sample_chain_map[res.parent.id], res.id[1]))
|
|
|
| dists = scipy.spatial.distance.cdist(coords, coords)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| for i in range(len(dists)):
|
| dists[i][max(i - 100, 0):] = np.inf
|
|
|
| clashes = np.argwhere(dists < 3)
|
| print("there are", len(clashes), "clashes")
|
|
|
|
|
|
|
|
|
| close_residues = np.argwhere(dists < INTERFACE_MIN_ATOM_DIST)
|
| inter_close_residues = [i for i in close_residues if identifiers[i[0]][0] != identifiers[i[1]][0]]
|
|
|
| interface_ident_residues = set()
|
| for i, j in inter_close_residues:
|
| identifier1, identifier2 = identifiers[i], identifiers[j]
|
| ident_chain1, ident_chain2 = ident_chains[identifier1[0]], ident_chains[identifier2[0]]
|
| ident_identifier1 = (ident_chain1, identifier1[1])
|
| ident_identifier2 = (ident_chain2, identifier2[1])
|
| if ident_identifier1 > ident_identifier2:
|
| ident_identifier1, ident_identifier2 = ident_identifier2, ident_identifier1
|
| interface_ident_residues.add((ident_identifier1, ident_identifier2))
|
|
|
| print("there are", len(inter_close_residues), "interface pairs and", len(interface_ident_residues),
|
| "interface ident pairs")
|
|
|
| return list(interface_ident_residues)
|
|
|
|
|
| def get_contacts_fast(pdb_path: str, ident_chains: Dict[str, str], sample_chain_map=None, offset: int = 0,
|
| bfactor_threshold=0):
|
| pdb_parser = Bio.PDB.PDBParser(QUIET=True)
|
| pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
|
| pdb_model = next(iter(pdb_struct))
|
|
|
| coords = []
|
| identifiers: List[Tuple[str, int]] = []
|
| for res in pdb_model.get_residues():
|
| if "CA" in res:
|
| if res["CA"].get_bfactor() < bfactor_threshold:
|
| continue
|
| for atom in res:
|
|
|
|
|
|
|
| coords.append(atom.get_coord())
|
| if sample_chain_map is None:
|
|
|
| identifiers.append((res.parent.id, res.id[1] + offset))
|
| else:
|
|
|
| identifiers.append((sample_chain_map[res.parent.id], res.id[1]))
|
|
|
| tree = KDTree(coords)
|
| close_residues = tree.query_pairs(INTERFACE_MIN_ATOM_DIST, p=2)
|
| inter_close_residues = [i for i in close_residues if identifiers[i[0]][0] != identifiers[i[1]][0]]
|
|
|
| interface_ident_residues = set()
|
| for i, j in inter_close_residues:
|
| identifier1, identifier2 = identifiers[i], identifiers[j]
|
| ident_chain1, ident_chain2 = ident_chains[identifier1[0]], ident_chains[identifier2[0]]
|
| ident_identifier1 = (ident_chain1, identifier1[1])
|
| ident_identifier2 = (ident_chain2, identifier2[1])
|
| if ident_identifier1 > ident_identifier2:
|
| ident_identifier1, ident_identifier2 = ident_identifier2, ident_identifier1
|
| interface_ident_residues.add((ident_identifier1, ident_identifier2))
|
|
|
| print("there are", len(inter_close_residues), "interface pairs and", len(interface_ident_residues),
|
| "interface ident pairs")
|
|
|
| return list(interface_ident_residues)
|
|
|
|
|
| def get_all_residues(pdb_path: str):
|
| pdb_parser = Bio.PDB.PDBParser(QUIET=True)
|
| pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
|
| pdb_model = next(iter(pdb_struct))
|
|
|
| identifiers: List[Tuple[str, int]] = []
|
| for res in pdb_model.get_residues():
|
| if "CA" in res:
|
| identifiers.append((res.parent.id, res.id[1]))
|
|
|
| return identifiers
|
|
|
|
|
| def get_contact_score(target_pdb_path: str, sample_pdb_path: str, sample_chain_map: Dict[str, str] = None,
|
| sample_offset: int = 0):
|
|
|
| ident_chains = get_ident_chain_map_from_complex(target_pdb_path)
|
|
|
|
|
|
|
|
|
|
|
|
|
| target_contacts = get_contacts_fast(target_pdb_path, ident_chains)
|
| sample_contacts = get_contacts_fast(sample_pdb_path, ident_chains, sample_chain_map, sample_offset,
|
| bfactor_threshold=50)
|
|
|
| all_target_residues = get_all_residues(target_pdb_path)
|
| sample_contacts = [contact for contact in sample_contacts
|
| if contact[0] in all_target_residues and contact[1] in all_target_residues]
|
| print("after filtering, there are", len(sample_contacts), "sample contacts")
|
|
|
| print(sorted(target_contacts)[:100])
|
| print(sorted(sample_contacts)[:100])
|
|
|
| precision = len(set(sample_contacts).intersection(target_contacts)) / len(sample_contacts)
|
| recall = len(set(sample_contacts).intersection(target_contacts)) / len(target_contacts)
|
|
|
| f1_score = 2 * precision * recall / (precision + recall)
|
|
|
| return f1_score
|
|
|
|
|
| def calc_scores():
|
| target_folder = "input_complexes"
|
|
|
| target_paths = {}
|
| for target_pdb_name in os.listdir(target_folder):
|
| if not target_pdb_name.endswith(".pdb"):
|
| continue
|
| target_paths[target_pdb_name[:4]] = os.path.join(target_folder, target_pdb_name)
|
|
|
|
|
| pdb_id_to_sample_chain_map = {}
|
| sample_folder = "top1_results/benchmark2"
|
| sample_paths = {}
|
| for sample_pdb_name in os.listdir(sample_folder):
|
| if not sample_pdb_name.endswith(".pdb"):
|
| continue
|
| sample_pdb_path = os.path.join(sample_folder, sample_pdb_name)
|
| sample_paths[sample_pdb_name[:4]] = sample_pdb_path
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| ics_scores = {}
|
| for pdb_id, target_pdb_path in target_paths.items():
|
| print("calculating score for", pdb_id, "...")
|
| if pdb_id not in sample_paths:
|
| ics_scores[pdb_id] = 0
|
| continue
|
| sample_pdb_path = sample_paths[pdb_id]
|
|
|
| if pdb_id in pdb_id_to_sample_chain_map:
|
| ics_scores[pdb_id] = get_contact_score(target_pdb_path, sample_pdb_path, pdb_id_to_sample_chain_map[pdb_id])
|
| else:
|
| ics_scores[pdb_id] = get_contact_score(target_pdb_path, sample_pdb_path, sample_offset=1)
|
| print(pdb_id, ics_scores[pdb_id])
|
| print(ics_scores)
|
|
|
|
|
| if __name__ == '__main__':
|
| calc_scores()
|
| exit(0)
|
|
|
|
|
