# 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}')