import copy import warnings from pathlib import Path import numpy as np import torch from Bio.PDB import PDBParser from rdkit import Chem from rdkit.Chem.rdchem import BondType as BT from rdkit.Chem import AllChem, GetPeriodicTable, RemoveHs from rdkit.Geometry import Point3D from torch import cdist from torch_cluster import knn_graph import prody as pr import torch.nn.functional as F from datasets.conformer_matching import get_torsion_angles, optimize_rotatable_bonds from datasets.constants import aa_short2long, atom_order, three_to_one from datasets.parse_chi import get_chi_angles, get_coords, aa_idx2aa_short, get_onehot_sequence from utils.torsion import get_transformation_mask from utils.logging_utils import get_logger logger = get_logger() periodic_table = GetPeriodicTable() allowable_features = { 'possible_atomic_num_list': list(range(1, 119)) + ['misc'], 'possible_chirality_list': [ 'CHI_UNSPECIFIED', 'CHI_TETRAHEDRAL_CW', 'CHI_TETRAHEDRAL_CCW', 'CHI_OTHER' ], 'possible_degree_list': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 'misc'], 'possible_numring_list': [0, 1, 2, 3, 4, 5, 6, 'misc'], 'possible_implicit_valence_list': [0, 1, 2, 3, 4, 5, 6, 'misc'], 'possible_formal_charge_list': [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 'misc'], 'possible_numH_list': [0, 1, 2, 3, 4, 5, 6, 7, 8, 'misc'], 'possible_number_radical_e_list': [0, 1, 2, 3, 4, 'misc'], 'possible_hybridization_list': [ 'SP', 'SP2', 'SP3', 'SP3D', 'SP3D2', 'misc' ], 'possible_is_aromatic_list': [False, True], 'possible_is_in_ring3_list': [False, True], 'possible_is_in_ring4_list': [False, True], 'possible_is_in_ring5_list': [False, True], 'possible_is_in_ring6_list': [False, True], 'possible_is_in_ring7_list': [False, True], 'possible_is_in_ring8_list': [False, True], 'possible_amino_acids': ['ALA', 'ARG', 'ASN', 'ASP', 'CYS', 'GLN', 'GLU', 'GLY', 'HIS', 'ILE', 'LEU', 'LYS', 'MET', 'PHE', 'PRO', 'SER', 'THR', 'TRP', 'TYR', 'VAL', 'HIP', 'HIE', 'TPO', 'HID', 'LEV', 'MEU', 'PTR', 'GLV', 'CYT', 'SEP', 'HIZ', 'CYM', 'GLM', 'ASQ', 'TYS', 'CYX', 'GLZ', 'misc'], 'possible_atom_type_2': ['C*', 'CA', 'CB', 'CD', 'CE', 'CG', 'CH', 'CZ', 'N*', 'ND', 'NE', 'NH', 'NZ', 'O*', 'OD', 'OE', 'OG', 'OH', 'OX', 'S*', 'SD', 'SG', 'misc'], 'possible_atom_type_3': ['C', 'CA', 'CB', 'CD', 'CD1', 'CD2', 'CE', 'CE1', 'CE2', 'CE3', 'CG', 'CG1', 'CG2', 'CH2', 'CZ', 'CZ2', 'CZ3', 'N', 'ND1', 'ND2', 'NE', 'NE1', 'NE2', 'NH1', 'NH2', 'NZ', 'O', 'OD1', 'OD2', 'OE1', 'OE2', 'OG', 'OG1', 'OH', 'OXT', 'SD', 'SG', 'misc'], } bonds = {BT.SINGLE: 0, BT.DOUBLE: 1, BT.TRIPLE: 2, BT.AROMATIC: 3} lig_feature_dims = (list(map(len, [ allowable_features['possible_atomic_num_list'], allowable_features['possible_chirality_list'], allowable_features['possible_degree_list'], allowable_features['possible_formal_charge_list'], allowable_features['possible_implicit_valence_list'], allowable_features['possible_numH_list'], allowable_features['possible_number_radical_e_list'], allowable_features['possible_hybridization_list'], allowable_features['possible_is_aromatic_list'], allowable_features['possible_numring_list'], allowable_features['possible_is_in_ring3_list'], allowable_features['possible_is_in_ring4_list'], allowable_features['possible_is_in_ring5_list'], allowable_features['possible_is_in_ring6_list'], allowable_features['possible_is_in_ring7_list'], allowable_features['possible_is_in_ring8_list'], ])), 0) # number of scalar features rec_atom_feature_dims = (list(map(len, [ allowable_features['possible_amino_acids'], allowable_features['possible_atomic_num_list'], allowable_features['possible_atom_type_2'], allowable_features['possible_atom_type_3'], ])), 0) rec_residue_feature_dims = (list(map(len, [ allowable_features['possible_amino_acids'] ])), 0) def lig_atom_featurizer(mol): ringinfo = mol.GetRingInfo() atom_features_list = [] for idx, atom in enumerate(mol.GetAtoms()): chiral_tag = str(atom.GetChiralTag()) if chiral_tag in ['CHI_SQUAREPLANAR', 'CHI_TRIGONALBIPYRAMIDAL', 'CHI_OCTAHEDRAL']: chiral_tag = 'CHI_OTHER' atom_features_list.append([ safe_index(allowable_features['possible_atomic_num_list'], atom.GetAtomicNum()), allowable_features['possible_chirality_list'].index(str(chiral_tag)), safe_index(allowable_features['possible_degree_list'], atom.GetTotalDegree()), safe_index(allowable_features['possible_formal_charge_list'], atom.GetFormalCharge()), safe_index(allowable_features['possible_implicit_valence_list'], atom.GetImplicitValence()), safe_index(allowable_features['possible_numH_list'], atom.GetTotalNumHs()), safe_index(allowable_features['possible_number_radical_e_list'], atom.GetNumRadicalElectrons()), safe_index(allowable_features['possible_hybridization_list'], str(atom.GetHybridization())), allowable_features['possible_is_aromatic_list'].index(atom.GetIsAromatic()), safe_index(allowable_features['possible_numring_list'], ringinfo.NumAtomRings(idx)), allowable_features['possible_is_in_ring3_list'].index(ringinfo.IsAtomInRingOfSize(idx, 3)), allowable_features['possible_is_in_ring4_list'].index(ringinfo.IsAtomInRingOfSize(idx, 4)), allowable_features['possible_is_in_ring5_list'].index(ringinfo.IsAtomInRingOfSize(idx, 5)), allowable_features['possible_is_in_ring6_list'].index(ringinfo.IsAtomInRingOfSize(idx, 6)), allowable_features['possible_is_in_ring7_list'].index(ringinfo.IsAtomInRingOfSize(idx, 7)), allowable_features['possible_is_in_ring8_list'].index(ringinfo.IsAtomInRingOfSize(idx, 8)), #g_charge if not np.isnan(g_charge) and not np.isinf(g_charge) else 0. ]) return torch.tensor(atom_features_list) def safe_index(l, e): """ Return index of element e in list l. If e is not present, return the last index """ try: return l.index(e) except: return len(l) - 1 def moad_extract_receptor_structure(path, complex_graph, neighbor_cutoff=20, max_neighbors=None, sequences_to_embeddings=None, knn_only_graph=False, lm_embeddings=None, all_atoms=False, atom_cutoff=None, atom_max_neighbors=None): # load the entire pdb file pdb = pr.parsePDB(str(path)) seq = pdb.ca.getSequence() coords = get_coords(pdb) one_hot = get_onehot_sequence(seq) chain_ids = np.zeros(len(one_hot)) res_chain_ids = pdb.ca.getChids() res_seg_ids = pdb.ca.getSegnames() res_chain_ids = np.asarray([s + c for s, c in zip(res_seg_ids, res_chain_ids)]) ids = np.unique(res_chain_ids) sequences = [] lm_embeddings = lm_embeddings if sequences_to_embeddings is None else [] for i, id in enumerate(ids): chain_ids[res_chain_ids == id] = i s = np.argmax(one_hot[res_chain_ids == id], axis=1) s = ''.join([aa_idx2aa_short[aa_idx] for aa_idx in s]) sequences.append(s) if sequences_to_embeddings is not None: lm_embeddings.append(sequences_to_embeddings[s]) complex_graph['receptor'].sequence = sequences complex_graph['receptor'].chain_ids = torch.from_numpy(np.asarray(chain_ids)).long() new_extract_receptor_structure(seq, coords, complex_graph, neighbor_cutoff=neighbor_cutoff, max_neighbors=max_neighbors, lm_embeddings=lm_embeddings, knn_only_graph=knn_only_graph, all_atoms=all_atoms, atom_cutoff=atom_cutoff, atom_max_neighbors=atom_max_neighbors) def new_extract_receptor_structure(seq, all_coords, complex_graph, neighbor_cutoff=20, max_neighbors=None, lm_embeddings=None, knn_only_graph=False, all_atoms=False, atom_cutoff=None, atom_max_neighbors=None): chi_angles, one_hot = get_chi_angles(all_coords, seq, return_onehot=True) n_rel_pos, c_rel_pos = all_coords[:, 0, :] - all_coords[:, 1, :], all_coords[:, 2, :] - all_coords[:, 1, :] side_chain_vecs = torch.from_numpy(np.concatenate([chi_angles / 360, n_rel_pos, c_rel_pos], axis=1)) # Build the k-NN graph coords = torch.tensor(all_coords[:, 1, :], dtype=torch.float) if len(coords) > 3000: raise ValueError(f'The receptor is too large {len(coords)}') if knn_only_graph: edge_index = knn_graph(coords, k=max_neighbors if max_neighbors else 32) else: distances = cdist(coords, coords) src_list = [] dst_list = [] for i in range(len(coords)): dst = list(np.where(distances[i, :] < neighbor_cutoff)[0]) dst.remove(i) max_neighbors = max_neighbors if max_neighbors else 1000 if max_neighbors != None and len(dst) > max_neighbors: dst = list(np.argsort(distances[i, :]))[1: max_neighbors + 1] if len(dst) == 0: dst = list(np.argsort(distances[i, :]))[1:2] # choose second because first is i itself print( f'The cutoff {neighbor_cutoff} was too small for one atom such that it had no neighbors. ' f'So we connected it to the closest other atom') assert i not in dst src = [i] * len(dst) src_list.extend(src) dst_list.extend(dst) edge_index = torch.from_numpy(np.asarray([dst_list, src_list])) res_names_list = [aa_short2long[seq[i]] if seq[i] in aa_short2long else 'misc' for i in range(len(seq))] feature_list = [[safe_index(allowable_features['possible_amino_acids'], res)] for res in res_names_list] node_feat = torch.tensor(feature_list, dtype=torch.float32) lm_embeddings = torch.tensor(np.concatenate(lm_embeddings, axis=0)) if lm_embeddings is not None else None complex_graph['receptor'].x = torch.cat([node_feat, lm_embeddings], axis=1) if lm_embeddings is not None else node_feat complex_graph['receptor'].pos = coords complex_graph['receptor'].side_chain_vecs = side_chain_vecs.float() complex_graph['receptor', 'rec_contact', 'receptor'].edge_index = edge_index if all_atoms: atom_coords = all_coords.reshape(-1, 3) atom_coords = torch.from_numpy(atom_coords[~np.any(np.isnan(atom_coords), axis=1)]).float() if knn_only_graph: atoms_edge_index = knn_graph(atom_coords, k=atom_max_neighbors if atom_max_neighbors else 1000) else: atoms_distances = cdist(atom_coords, atom_coords) atom_src_list = [] atom_dst_list = [] for i in range(len(atom_coords)): dst = list(np.where(atoms_distances[i, :] < atom_cutoff)[0]) dst.remove(i) max_neighbors = atom_max_neighbors if atom_max_neighbors else 1000 if max_neighbors != None and len(dst) > max_neighbors: dst = list(np.argsort(atoms_distances[i, :]))[1: max_neighbors + 1] if len(dst) == 0: dst = list(np.argsort(atoms_distances[i, :]))[1:2] # choose second because first is i itself print( f'The atom_cutoff {atom_cutoff} was too small for one atom such that it had no neighbors. ' f'So we connected it to the closest other atom') assert i not in dst src = [i] * len(dst) atom_src_list.extend(src) atom_dst_list.extend(dst) atoms_edge_index = torch.from_numpy(np.asarray([atom_dst_list, atom_src_list])) feats = [get_moad_atom_feats(res, all_coords[i]) for i, res in enumerate(seq)] atom_feat = torch.from_numpy(np.concatenate(feats, axis=0)).float() c_alpha_idx = np.concatenate([np.zeros(len(f)) + i for i, f in enumerate(feats)]) np_array = np.stack([np.arange(len(atom_feat)), c_alpha_idx]) atom_res_edge_index = torch.from_numpy(np_array).long() complex_graph['atom'].x = atom_feat complex_graph['atom'].pos = atom_coords assert len(complex_graph['atom'].x) == len(complex_graph['atom'].pos) complex_graph['atom', 'atom_contact', 'atom'].edge_index = atoms_edge_index complex_graph['atom', 'atom_rec_contact', 'receptor'].edge_index = atom_res_edge_index return def get_moad_atom_feats(res, coords): feats = [] res_long = aa_short2long[res] res_order = atom_order[res] for i, c in enumerate(coords): if np.any(np.isnan(c)): continue atom_feats = [] if res == '-': atom_feats = [safe_index(allowable_features['possible_amino_acids'], 'misc'), safe_index(allowable_features['possible_atomic_num_list'], 'misc'), safe_index(allowable_features['possible_atom_type_2'], 'misc'), safe_index(allowable_features['possible_atom_type_3'], 'misc')] else: atom_feats.append(safe_index(allowable_features['possible_amino_acids'], res_long)) if i >= len(res_order): atom_feats.extend([safe_index(allowable_features['possible_atomic_num_list'], 'misc'), safe_index(allowable_features['possible_atom_type_2'], 'misc'), safe_index(allowable_features['possible_atom_type_3'], 'misc')]) else: atom_name = res_order[i] try: atomic_num = periodic_table.GetAtomicNumber(atom_name[:1]) except: print("element", res_order[i][:1], 'not found') atomic_num = -1 atom_feats.extend([safe_index(allowable_features['possible_atomic_num_list'], atomic_num), safe_index(allowable_features['possible_atom_type_2'], (atom_name + '*')[:2]), safe_index(allowable_features['possible_atom_type_3'], atom_name)]) feats.append(atom_feats) feats = np.asarray(feats) return feats def get_lig_graph(mol, complex_graph): atom_feats = lig_atom_featurizer(mol) row, col, edge_type = [], [], [] for bond in mol.GetBonds(): start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx() row += [start, end] col += [end, start] edge_type += 2 * [bonds[bond.GetBondType()]] if bond.GetBondType() != BT.UNSPECIFIED else [0, 0] edge_index = torch.tensor([row, col], dtype=torch.long) edge_type = torch.tensor(edge_type, dtype=torch.long) edge_attr = F.one_hot(edge_type, num_classes=len(bonds)).to(torch.float) complex_graph['ligand'].x = atom_feats complex_graph['ligand', 'lig_bond', 'ligand'].edge_index = edge_index complex_graph['ligand', 'lig_bond', 'ligand'].edge_attr = edge_attr if mol.GetNumConformers() > 0: lig_coords = torch.from_numpy(mol.GetConformer().GetPositions()).float() complex_graph['ligand'].pos = lig_coords return def generate_conformer(mol): ps = AllChem.ETKDGv2() failures, id = 0, -1 while failures < 3 and id == -1: if failures > 0: logger.debug(f'rdkit coords could not be generated. trying again {failures}.') id = AllChem.EmbedMolecule(mol, ps) failures += 1 if id == -1: logger.info('rdkit coords could not be generated without using random coords. using random coords now.') ps.useRandomCoords = True AllChem.EmbedMolecule(mol, ps) AllChem.MMFFOptimizeMolecule(mol, confId=0) return True #else: # AllChem.MMFFOptimizeMolecule(mol, confId=0) return False def get_lig_graph_with_matching(mol_, complex_graph, popsize, maxiter, matching, keep_original, num_conformers, remove_hs, tries=10, skip_matching=False): if matching: mol_maybe_noh = copy.deepcopy(mol_) if remove_hs: mol_maybe_noh = RemoveHs(mol_maybe_noh, sanitize=True) mol_maybe_noh = AllChem.RemoveAllHs(mol_maybe_noh) if keep_original: positions = [] for conf in mol_maybe_noh.GetConformers(): positions.append(conf.GetPositions()) complex_graph['ligand'].orig_pos = np.asarray(positions) if len(positions) > 1 else positions[0] # rotatable_bonds = get_torsion_angles(mol_maybe_noh) _tmp = copy.deepcopy(mol_) if remove_hs: _tmp = RemoveHs(_tmp, sanitize=True) _tmp = AllChem.RemoveAllHs(_tmp) rotatable_bonds = get_torsion_angles(_tmp) for i in range(num_conformers): mols, rmsds = [], [] for _ in range(tries): mol_rdkit = copy.deepcopy(mol_) mol_rdkit.RemoveAllConformers() mol_rdkit = AllChem.AddHs(mol_rdkit) generate_conformer(mol_rdkit) if remove_hs: mol_rdkit = RemoveHs(mol_rdkit, sanitize=True) mol_rdkit = AllChem.RemoveAllHs(mol_rdkit) mol = AllChem.RemoveAllHs(copy.deepcopy(mol_maybe_noh)) if rotatable_bonds and not skip_matching: optimize_rotatable_bonds(mol_rdkit, mol, rotatable_bonds, popsize=popsize, maxiter=maxiter) mol.AddConformer(mol_rdkit.GetConformer()) rms_list = [] AllChem.AlignMolConformers(mol, RMSlist=rms_list) mol_rdkit.RemoveAllConformers() mol_rdkit.AddConformer(mol.GetConformers()[1]) mols.append(mol_rdkit) rmsds.append(rms_list[0]) # select molecule with lowest rmsd #print("mean std min max", np.mean(rmsds), np.std(rmsds), np.min(rmsds), np.max(rmsds)) mol_rdkit = mols[np.argmin(rmsds)] if i == 0: complex_graph.rmsd_matching = min(rmsds) get_lig_graph(mol_rdkit, complex_graph) else: if torch.is_tensor(complex_graph['ligand'].pos): complex_graph['ligand'].pos = [complex_graph['ligand'].pos] complex_graph['ligand'].pos.append(torch.from_numpy(mol_rdkit.GetConformer().GetPositions()).float()) else: # no matching complex_graph.rmsd_matching = 0 if remove_hs: mol_ = RemoveHs(mol_) get_lig_graph(mol_, complex_graph) edge_mask, mask_rotate = get_transformation_mask(complex_graph) complex_graph['ligand'].edge_mask = torch.tensor(edge_mask) complex_graph['ligand'].mask_rotate = mask_rotate return def get_rec_misc_atom_feat(bio_atom=None, atom_name=None, element=None, get_misc_features=False): if get_misc_features: return [safe_index(allowable_features['possible_amino_acids'], 'misc'), safe_index(allowable_features['possible_atomic_num_list'], 'misc'), safe_index(allowable_features['possible_atom_type_2'], 'misc'), safe_index(allowable_features['possible_atom_type_3'], 'misc')] if atom_name is not None: atom_name = atom_name else: atom_name = bio_atom.name if element is not None: element = element else: element = bio_atom.element if element == 'CD': element = 'C' assert not element == '' try: atomic_num = periodic_table.GetAtomicNumber(element.lower().capitalize()) except: atomic_num = -1 atom_feat = [safe_index(allowable_features['possible_amino_acids'], bio_atom.get_parent().get_resname()), safe_index(allowable_features['possible_atomic_num_list'], atomic_num), safe_index(allowable_features['possible_atom_type_2'], (atom_name + '*')[:2]), safe_index(allowable_features['possible_atom_type_3'], atom_name)] return atom_feat def write_mol_with_coords(mol, new_coords, path): w = Chem.SDWriter(path) conf = mol.GetConformer() for i in range(mol.GetNumAtoms()): x,y,z = new_coords.astype(np.double)[i] conf.SetAtomPosition(i,Point3D(x,y,z)) w.write(mol) w.close() def create_mol_with_coords(mol, new_coords, path=None): conf = mol.GetConformer() for i in range(mol.GetNumAtoms()): x, y, z = new_coords[i] conf.SetAtomPosition(i, Point3D(float(x), float(y), float(z))) if path: w = Chem.SDWriter(path) w.write(mol) w.close() return mol def read_molecule(ligand_description, sanitize=False, calc_charges=False, remove_hs=False, remove_confs=False): mol = None # Check if ligand_description is a path to a file if Path(ligand_description).is_absolute() or len(Path(ligand_description).parts) > 1: path = Path(ligand_description) if path.is_file(): match path.suffix: case '.mol': mol = Chem.MolFromMolFile(str(path), sanitize=False, removeHs=True) case '.mol2': mol = Chem.MolFromMol2File(str(path), sanitize=False, removeHs=False) case '.sdf': supplier = Chem.SDMolSupplier(str(path), sanitize=False, removeHs=False) mol = supplier[0] case '.pdbqt': with open(path) as file: pdbqt_data = file.readlines() pdb_block = '' for line in pdbqt_data: pdb_block += '{}\n'.format(line[:66]) mol = Chem.MolFromPDBBlock(pdb_block, sanitize=False, removeHs=False) case '.pdb': mol = Chem.MolFromPDBFile(str(path), sanitize=False, removeHs=False) case _: logger.warning(f'Expect the format of the molecule file to be ' f'one of .mol2, .sdf, .pdbqt and .pdb, got {ligand_description}') else: raise FileNotFoundError(f'File {ligand_description} not found.') else: mol = Chem.MolFromSmiles(ligand_description) # No need to remove conformers if the molecule is not read from a file remove_confs = False if mol is not None: try: if sanitize or calc_charges: Chem.SanitizeMol(mol) if calc_charges: # Compute Gasteiger charges on the molecule. try: AllChem.ComputeGasteigerCharges(mol) except: warnings.warn('Unable to compute charges for the molecule.') if remove_hs: mol = Chem.RemoveHs(mol, sanitize=sanitize) if remove_confs: mol.RemoveAllConformers() except Exception as e: # Print stacktrace import traceback msg = traceback.format_exc() logger.warning(f"Failed to process molecule: {ligand_description}\n{msg}") return None return mol