# Copyright (c) DP Techonology, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import numpy as np from rdkit import Chem from rdkit.Chem import AllChem from rdkit import RDLogger RDLogger.DisableLog("rdApp.*") import warnings warnings.filterwarnings(action="ignore") from rdkit.Chem import rdMolTransforms import copy import lmdb import pickle import pandas as pd def get_torsions(m, removeHs=True): if removeHs: m = Chem.RemoveHs(m) torsionList = [] torsionSmarts = "[!$(*#*)&!D1]-&!@[!$(*#*)&!D1]" torsionQuery = Chem.MolFromSmarts(torsionSmarts) matches = m.GetSubstructMatches(torsionQuery) for match in matches: idx2 = match[0] idx3 = match[1] bond = m.GetBondBetweenAtoms(idx2, idx3) jAtom = m.GetAtomWithIdx(idx2) kAtom = m.GetAtomWithIdx(idx3) for b1 in jAtom.GetBonds(): if b1.GetIdx() == bond.GetIdx(): continue idx1 = b1.GetOtherAtomIdx(idx2) for b2 in kAtom.GetBonds(): if (b2.GetIdx() == bond.GetIdx()) or (b2.GetIdx() == b1.GetIdx()): continue idx4 = b2.GetOtherAtomIdx(idx3) # skip 3-membered rings if idx4 == idx1: continue # skip torsions that include hydrogens if (m.GetAtomWithIdx(idx1).GetAtomicNum() == 1) or ( m.GetAtomWithIdx(idx4).GetAtomicNum() == 1 ): continue if m.GetAtomWithIdx(idx4).IsInRing(): torsionList.append((idx4, idx3, idx2, idx1)) break else: torsionList.append((idx1, idx2, idx3, idx4)) break break return torsionList def SetDihedral(conf, atom_idx, new_vale): rdMolTransforms.SetDihedralRad( conf, atom_idx[0], atom_idx[1], atom_idx[2], atom_idx[3], new_vale ) def single_conf_gen_bonds(tgt_mol, num_confs=1000, seed=42, removeHs=True): mol = copy.deepcopy(tgt_mol) mol = Chem.AddHs(mol) allconformers = AllChem.EmbedMultipleConfs( mol, numConfs=num_confs, randomSeed=seed, clearConfs=True ) if removeHs: mol = Chem.RemoveHs(mol) rotable_bonds = get_torsions(mol, removeHs=removeHs) for i in range(len(allconformers)): np.random.seed(i) values = 3.1415926 * 2 * np.random.rand(len(rotable_bonds)) for idx in range(len(rotable_bonds)): SetDihedral(mol.GetConformers()[i], rotable_bonds[idx], values[idx]) Chem.rdMolTransforms.CanonicalizeConformer(mol.GetConformers()[i]) return mol def load_lmdb_data(lmdb_path, key): env = lmdb.open( lmdb_path, subdir=False, readonly=True, lock=False, readahead=False, meminit=False, max_readers=256, ) txn = env.begin() _keys = list(txn.cursor().iternext(values=False)) collects = [] for idx in range(len(_keys)): datapoint_pickled = txn.get(f"{idx}".encode("ascii")) data = pickle.loads(datapoint_pickled) collects.append(data[key]) return collects def docking_data_pre(raw_data_path, predict_path): mol_list = load_lmdb_data(raw_data_path, "mol_list") mol_list = [Chem.RemoveHs(mol) for items in mol_list for mol in items] predict = pd.read_pickle(predict_path) ( smi_list, pocket_list, pocket_coords_list, distance_predict_list, holo_distance_predict_list, holo_coords_list, holo_center_coords_list, ) = ([], [], [], [], [], [], []) for batch in predict: sz = batch["atoms"].size(0) for i in range(sz): smi_list.append(batch["smi_name"][i]) pocket_list.append(batch["pocket_name"][i]) distance_predict = batch["cross_distance_predict"][i] token_mask = batch["atoms"][i] > 2 pocket_token_mask = batch["pocket_atoms"][i] > 2 distance_predict = distance_predict[token_mask][:, pocket_token_mask] pocket_coords = batch["pocket_coordinates"][i] pocket_coords = pocket_coords[pocket_token_mask, :] holo_distance_predict = batch["holo_distance_predict"][i] holo_distance_predict = holo_distance_predict[token_mask][:, token_mask] holo_coordinates = batch["holo_coordinates"][i] holo_coordinates = holo_coordinates[token_mask, :] holo_center_coordinates = batch["holo_center_coordinates"][i][:3] pocket_coords = pocket_coords.numpy().astype(np.float32) distance_predict = distance_predict.numpy().astype(np.float32) holo_distance_predict = holo_distance_predict.numpy().astype(np.float32) holo_coords = holo_coordinates.numpy().astype(np.float32) pocket_coords_list.append(pocket_coords) distance_predict_list.append(distance_predict) holo_distance_predict_list.append(holo_distance_predict) holo_coords_list.append(holo_coords) holo_center_coords_list.append(holo_center_coordinates) return ( mol_list, smi_list, pocket_list, pocket_coords_list, distance_predict_list, holo_distance_predict_list, holo_coords_list, holo_center_coords_list, ) def ensemble_iterations( mol_list, smi_list, pocket_list, pocket_coords_list, distance_predict_list, holo_distance_predict_list, holo_coords_list, holo_center_coords_list, tta_times=10, ): sz = len(mol_list) for i in range(sz // tta_times): start_idx, end_idx = i * tta_times, (i + 1) * tta_times distance_predict_tta = distance_predict_list[start_idx:end_idx] holo_distance_predict_tta = holo_distance_predict_list[start_idx:end_idx] mol = copy.deepcopy(mol_list[start_idx]) rdkit_mol = single_conf_gen_bonds( mol, num_confs=tta_times, seed=42, removeHs=True ) sz = len(rdkit_mol.GetConformers()) initial_coords_list = [ rdkit_mol.GetConformers()[i].GetPositions().astype(np.float32) for i in range(sz) ] yield [ initial_coords_list, mol, smi_list[start_idx], pocket_list[start_idx], pocket_coords_list[start_idx], distance_predict_tta, holo_distance_predict_tta, holo_coords_list[start_idx], holo_center_coords_list[start_idx], ] def rmsd_func(holo_coords, predict_coords): if predict_coords is not np.nan: sz = holo_coords.shape rmsd = np.sqrt(np.sum((predict_coords - holo_coords) ** 2) / sz[0]) return rmsd return 1000.0 def print_results(rmsd_results): print("RMSD < 1.0 : ", np.mean(rmsd_results < 1.0)) print("RMSD < 1.5 : ", np.mean(rmsd_results < 1.5)) print("RMSD < 2.0 : ", np.mean(rmsd_results < 2.0)) print("RMSD < 3.0 : ", np.mean(rmsd_results < 3.0)) print("RMSD < 5.0 : ", np.mean(rmsd_results < 5.0)) print("avg RMSD : ", np.mean(rmsd_results))