# 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))
Datasets
Models
