# This file needs to be ran in the TANKBind repository together with baseline_run_tankbind_parallel.sh
import sys
import time
from multiprocessing import Pool
import copy
import warnings
from argparse import ArgumentParser
from rdkit.Chem import AllChem, RemoveHs
from feature_utils import save_cleaned_protein, read_mol
from generation_utils import get_LAS_distance_constraint_mask, get_info_pred_distance, write_with_new_coords
import logging
from torch_geometric.loader import DataLoader
from tqdm import tqdm # pip install tqdm if fails.
from model import get_model
# from utils import *
import torch
from data import TankBind_prediction
import os
import numpy as np
import pandas as pd
import rdkit.Chem as Chem
from feature_utils import generate_sdf_from_smiles_using_rdkit
from feature_utils import get_protein_feature
from Bio.PDB import PDBParser
from feature_utils import extract_torchdrug_feature_from_mol
def read_strings_from_txt(path):
# every line will be one element of the returned list
with open(path) as file:
lines = file.readlines()
return [line.rstrip() for line in lines]
def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=False):
if molecule_file.endswith('.mol2'):
mol = Chem.MolFromMol2File(molecule_file, sanitize=False, removeHs=False)
elif molecule_file.endswith('.sdf'):
supplier = Chem.SDMolSupplier(molecule_file, sanitize=False, removeHs=False)
mol = supplier[0]
elif molecule_file.endswith('.pdbqt'):
with open(molecule_file) 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)
elif molecule_file.endswith('.pdb'):
mol = Chem.MolFromPDBFile(molecule_file, sanitize=False, removeHs=False)
else:
return ValueError('Expect the format of the molecule_file to be '
'one of .mol2, .sdf, .pdbqt and .pdb, got {}'.format(molecule_file))
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)
except:
return None
return mol
def parallel_save_prediction(arguments):
dataset, y_pred_list, chosen,rdkit_mol_path, result_folder, name = arguments
for idx, line in chosen.iterrows():
pocket_name = line['pocket_name']
compound_name = line['compound_name']
ligandName = compound_name.split("_")[1]
dataset_index = line['dataset_index']
coords = dataset[dataset_index].coords.to('cpu')
protein_nodes_xyz = dataset[dataset_index].node_xyz.to('cpu')
n_compound = coords.shape[0]
n_protein = protein_nodes_xyz.shape[0]
y_pred = y_pred_list[dataset_index].reshape(n_protein, n_compound).to('cpu')
compound_pair_dis_constraint = torch.cdist(coords, coords)
mol = Chem.MolFromMolFile(rdkit_mol_path)
LAS_distance_constraint_mask = get_LAS_distance_constraint_mask(mol).bool()
pred_dist_info = get_info_pred_distance(coords, y_pred, protein_nodes_xyz, compound_pair_dis_constraint,
LAS_distance_constraint_mask=LAS_distance_constraint_mask,
n_repeat=1, show_progress=False)
toFile = f'{result_folder}/{name}_tankbind_chosen.sdf'
new_coords = pred_dist_info.sort_values("loss")['coords'].iloc[0].astype(np.double)
write_with_new_coords(mol, new_coords, toFile)
if __name__ == '__main__':
tankbind_src_folder = "../tankbind"
sys.path.insert(0, tankbind_src_folder)
torch.set_num_threads(16)
parser = ArgumentParser()
parser.add_argument('--data_dir', type=str, default='/Users/hstark/projects/ligbind/data/PDBBind_processed', help='')
parser.add_argument('--split_path', type=str, default='/Users/hstark/projects/ligbind/data/splits/timesplit_test', help='')
parser.add_argument('--prank_path', type=str, default='/Users/hstark/projects/p2rank_2.3/prank', help='')
parser.add_argument('--results_path', type=str, default='results/tankbind_results', help='')
parser.add_argument('--skip_existing', action='store_true', default=False, help='')
parser.add_argument('--skip_p2rank', action='store_true', default=False, help='')
parser.add_argument('--skip_multiple_pocket_outputs', action='store_true', default=False, help='')
parser.add_argument('--device', type=str, default='cpu', help='')
parser.add_argument('--num_workers', type=int, default=1, help='')
parser.add_argument('--parallel_id', type=int, default=0, help='')
parser.add_argument('--parallel_tot', type=int, default=1, help='')
args = parser.parse_args()
device = args.device
cache_path = "tankbind_cache"
os.makedirs(cache_path, exist_ok=True)
os.makedirs(args.results_path, exist_ok=True)
logging.basicConfig(level=logging.INFO)
model = get_model(0, logging, device)
# re-dock model
# modelFile = "../saved_models/re_dock.pt"
# self-dock model
modelFile = f"{tankbind_src_folder}/../saved_models/self_dock.pt"
model.load_state_dict(torch.load(modelFile, map_location=device))
_ = model.eval()
batch_size = 5
names = read_strings_from_txt(args.split_path)
if args.parallel_tot > 1:
size = len(names) // args.parallel_tot + 1
names = names[args.parallel_id*size:(args.parallel_id+1)*size]
rmsds = []
forward_pass_time = []
times_preprocess = []
times_inference = []
top_10_generation_time = []
top_1_generation_time = []
start_time = time.time()
if not args.skip_p2rank:
for name in names:
if args.skip_existing and os.path.exists(f'{args.results_path}/{name}/{name}_tankbind_1.sdf'): continue
print("Now processing: ", name)
protein_path = f'{args.data_dir}/{name}/{name}_protein_processed.pdb'
cleaned_protein_path = f"{cache_path}/{name}_protein_tankbind_cleaned.pdb" # if you change this you also need to change below
parser = PDBParser(QUIET=True)
s = parser.get_structure(name, protein_path)
c = s[0]
clean_res_list, ligand_list = save_cleaned_protein(c, cleaned_protein_path)
with open(f"{cache_path}/pdb_list_p2rank.txt", "w") as out:
for name in names:
out.write(f"{name}_protein_tankbind_cleaned.pdb\n")
cmd = f"bash {args.prank_path} predict {cache_path}/pdb_list_p2rank.txt -o {cache_path}/p2rank -threads 4"
os.system(cmd)
times_preprocess.append(time.time() - start_time)
p2_rank_time = time.time() - start_time
list_to_parallelize = []
for name in tqdm(names):
single_preprocess_time = time.time()
if args.skip_existing and os.path.exists(f'{args.results_path}/{name}/{name}_tankbind_1.sdf'): continue
print("Now processing: ", name)
protein_path = f'{args.data_dir}/{name}/{name}_protein_processed.pdb'
ligand_path = f"{args.data_dir}/{name}/{name}_ligand.sdf"
cleaned_protein_path = f"{cache_path}/{name}_protein_tankbind_cleaned.pdb" # if you change this you also need to change below
rdkit_mol_path = f"{cache_path}/{name}_rdkit_ligand.sdf"
parser = PDBParser(QUIET=True)
s = parser.get_structure(name, protein_path)
c = s[0]
clean_res_list, ligand_list = save_cleaned_protein(c, cleaned_protein_path)
lig, _ = read_mol(f"{args.data_dir}/{name}/{name}_ligand.sdf", f"{args.data_dir}/{name}/{name}_ligand.mol2")
lig = RemoveHs(lig)
smiles = Chem.MolToSmiles(lig)
generate_sdf_from_smiles_using_rdkit(smiles, rdkit_mol_path, shift_dis=0)
parser = PDBParser(QUIET=True)
s = parser.get_structure("x", cleaned_protein_path)
res_list = list(s.get_residues())
protein_dict = {}
protein_dict[name] = get_protein_feature(res_list)
compound_dict = {}
mol = Chem.MolFromMolFile(rdkit_mol_path)
compound_dict[name + f"_{name}" + "_rdkit"] = extract_torchdrug_feature_from_mol(mol, has_LAS_mask=True)
info = []
for compound_name in list(compound_dict.keys()):
# use protein center as the block center.
com = ",".join([str(a.round(3)) for a in protein_dict[name][0].mean(axis=0).numpy()])
info.append([name, compound_name, "protein_center", com])
p2rankFile = f"{cache_path}/p2rank/{name}_protein_tankbind_cleaned.pdb_predictions.csv"
pocket = pd.read_csv(p2rankFile)
pocket.columns = pocket.columns.str.strip()
pocket_coms = pocket[['center_x', 'center_y', 'center_z']].values
for ith_pocket, com in enumerate(pocket_coms):
com = ",".join([str(a.round(3)) for a in com])
info.append([name, compound_name, f"pocket_{ith_pocket + 1}", com])
info = pd.DataFrame(info, columns=['protein_name', 'compound_name', 'pocket_name', 'pocket_com'])
dataset_path = f"{cache_path}/{name}_dataset/"
os.system(f"rm -r {dataset_path}")
os.system(f"mkdir -p {dataset_path}")
dataset = TankBind_prediction(dataset_path, data=info, protein_dict=protein_dict, compound_dict=compound_dict)
# dataset = TankBind_prediction(dataset_path)
times_preprocess.append(time.time() - single_preprocess_time)
single_forward_pass_time = time.time()
data_loader = DataLoader(dataset, batch_size=batch_size, follow_batch=['x', 'y', 'compound_pair'], shuffle=False,
num_workers=0)
affinity_pred_list = []
y_pred_list = []
for data in tqdm(data_loader):
data = data.to(device)
y_pred, affinity_pred = model(data)
affinity_pred_list.append(affinity_pred.detach().cpu())
for i in range(data.y_batch.max() + 1):
y_pred_list.append((y_pred[data['y_batch'] == i]).detach().cpu())
affinity_pred_list = torch.cat(affinity_pred_list)
forward_pass_time.append(time.time() - single_forward_pass_time)
output_info = copy.deepcopy(dataset.data)
output_info['affinity'] = affinity_pred_list
output_info['dataset_index'] = range(len(output_info))
output_info_sorted = output_info.sort_values('affinity', ascending=False)
result_folder = f'{args.results_path}/{name}'
os.makedirs(result_folder, exist_ok=True)
output_info_sorted.to_csv(f"{result_folder}/output_info_sorted_by_affinity.csv")
if not args.skip_multiple_pocket_outputs:
for idx, (dataframe_idx, line) in enumerate(copy.deepcopy(output_info_sorted).iterrows()):
single_top10_generation_time = time.time()
pocket_name = line['pocket_name']
compound_name = line['compound_name']
ligandName = compound_name.split("_")[1]
coords = dataset[dataframe_idx].coords.to('cpu')
protein_nodes_xyz = dataset[dataframe_idx].node_xyz.to('cpu')
n_compound = coords.shape[0]
n_protein = protein_nodes_xyz.shape[0]
y_pred = y_pred_list[dataframe_idx].reshape(n_protein, n_compound).to('cpu')
y = dataset[dataframe_idx].dis_map.reshape(n_protein, n_compound).to('cpu')
compound_pair_dis_constraint = torch.cdist(coords, coords)
mol = Chem.MolFromMolFile(rdkit_mol_path)
LAS_distance_constraint_mask = get_LAS_distance_constraint_mask(mol).bool()
pred_dist_info = get_info_pred_distance(coords, y_pred, protein_nodes_xyz, compound_pair_dis_constraint,
LAS_distance_constraint_mask=LAS_distance_constraint_mask,
n_repeat=1, show_progress=False)
toFile = f'{result_folder}/{name}_tankbind_{idx}.sdf'
new_coords = pred_dist_info.sort_values("loss")['coords'].iloc[0].astype(np.double)
write_with_new_coords(mol, new_coords, toFile)
if idx < 10:
top_10_generation_time.append(time.time() - single_top10_generation_time)
if idx == 0:
top_1_generation_time.append(time.time() - single_top10_generation_time)
output_info_chosen = copy.deepcopy(dataset.data)
output_info_chosen['affinity'] = affinity_pred_list
output_info_chosen['dataset_index'] = range(len(output_info_chosen))
chosen = output_info_chosen.loc[
output_info_chosen.groupby(['protein_name', 'compound_name'], sort=False)['affinity'].agg(
'idxmax')].reset_index()
list_to_parallelize.append((dataset, y_pred_list, chosen, rdkit_mol_path, result_folder, name))
chosen_generation_start_time = time.time()
if args.num_workers > 1:
p = Pool(args.num_workers, maxtasksperchild=1)
p.__enter__()
with tqdm(total=len(list_to_parallelize), desc=f'running optimization {i}/{len(list_to_parallelize)}') as pbar:
map_fn = p.imap_unordered if args.num_workers > 1 else map
for t in map_fn(parallel_save_prediction, list_to_parallelize):
pbar.update()
if args.num_workers > 1: p.__exit__(None, None, None)
chosen_generation_time = time.time() - chosen_generation_start_time
"""
lig, _ = read_mol(f"{args.data_dir}/{name}/{name}_ligand.sdf", f"{args.data_dir}/{name}/{name}_ligand.mol2")
sm = Chem.MolToSmiles(lig)
m_order = list(lig.GetPropsAsDict(includePrivate=True, includeComputed=True)['_smilesAtomOutputOrder'])
lig = Chem.RenumberAtoms(lig, m_order)
lig = Chem.RemoveAllHs(lig)
lig = RemoveHs(lig)
true_ligand_pos = np.array(lig.GetConformer().GetPositions())
toFile = f'{result_folder}/{name}_tankbind_chosen.sdf'
mol_pred, _ = read_mol(toFile, None)
sm = Chem.MolToSmiles(mol_pred)
m_order = list(mol_pred.GetPropsAsDict(includePrivate=True, includeComputed=True)['_smilesAtomOutputOrder'])
mol_pred = Chem.RenumberAtoms(mol_pred, m_order)
mol_pred = RemoveHs(mol_pred)
mol_pred_pos = np.array(mol_pred.GetConformer().GetPositions())
rmsds.append(np.sqrt(((true_ligand_pos - mol_pred_pos) ** 2).sum(axis=1).mean(axis=0)))
print(np.sqrt(((true_ligand_pos - mol_pred_pos) ** 2).sum(axis=1).mean(axis=0)))
"""
forward_pass_time = np.array(forward_pass_time).sum()
times_preprocess = np.array(times_preprocess).sum()
times_inference = np.array(times_inference).sum()
top_10_generation_time = np.array(top_10_generation_time).sum()
top_1_generation_time = np.array(top_1_generation_time).sum()
rmsds = np.array(rmsds)
print(f'forward_pass_time: {forward_pass_time}')
print(f'times_preprocess: {times_preprocess}')
print(f'times_inference: {times_inference}')
print(f'top_10_generation_time: {top_10_generation_time}')
print(f'top_1_generation_time: {top_1_generation_time}')
print(f'chosen_generation_time: {chosen_generation_time}')
print(f'rmsds_below_2: {(100 * (rmsds < 2).sum() / len(rmsds))}')
print(f'p2rank Time: {p2_rank_time}')
print(
f'total_time: '
f'{forward_pass_time + times_preprocess + times_inference + top_10_generation_time + top_1_generation_time + p2_rank_time}')
with open(os.path.join(args.results_path, 'tankbind_log.log'), 'w') as file:
file.write(f'forward_pass_time: {forward_pass_time}')
file.write(f'times_preprocess: {times_preprocess}')
file.write(f'times_inference: {times_inference}')
file.write(f'top_10_generation_time: {top_10_generation_time}')
file.write(f'top_1_generation_time: {top_1_generation_time}')
file.write(f'rmsds_below_2: {(100 * (rmsds < 2).sum() / len(rmsds))}')
file.write(f'p2rank Time: {p2_rank_time}')
file.write(f'total_time: {forward_pass_time + times_preprocess + times_inference + top_10_generation_time + top_1_generation_time + p2_rank_time}')