| import copy | |
| import os | |
| import plotly.express as px | |
| import time | |
| from argparse import FileType, ArgumentParser | |
| import numpy as np | |
| import pandas as pd | |
| import wandb | |
| from biopandas.pdb import PandasPdb | |
| from rdkit import Chem | |
| from rdkit.Chem import RemoveHs | |
| from tqdm import tqdm | |
| from datasets.pdbbind import read_mol | |
| from datasets.process_mols import read_molecule, read_sdf_or_mol2 | |
| from utils.utils import read_strings_from_txt, get_symmetry_rmsd, remove_all_hs | |
| parser = ArgumentParser() | |
| parser.add_argument('--config', type=FileType(mode='r'), default=None) | |
| parser.add_argument('--run_name', type=str, default='tankbind', help='') | |
| parser.add_argument('--data_dir', type=str, default='data/PDBBind_processed', help='') | |
| parser.add_argument('--renumbered_atoms_dir', type=str, default='../TankBind/examples/tankbind_pdb/renumber_atom_index_same_as_smiles', help='') | |
| parser.add_argument('--results_path', type=str, default='results/tankbind_top5', help='Path to folder with trained model and hyperparameters') | |
| parser.add_argument('--project', type=str, default='ligbind_inf', help='') | |
| parser.add_argument('--wandb', action='store_true', default=True, help='') | |
| parser.add_argument('--num_predictions', type=int, default=5, help='') | |
| args = parser.parse_args() | |
| names = read_strings_from_txt(f'data/splits/timesplit_test') | |
| names_no_rec_overlap = read_strings_from_txt(f'data/splits/timesplit_test_no_rec_overlap') | |
| if args.wandb: | |
| wandb.init( | |
| entity='coarse-graining-mit', | |
| settings=wandb.Settings(start_method="fork"), | |
| project=args.project, | |
| name=args.run_name, | |
| config=args | |
| ) | |
| all_times = [] | |
| rmsds_list = [] | |
| unsym_rmsds_list = [] | |
| centroid_distances_list = [] | |
| min_cross_distances_list = [] | |
| min_self_distances_list = [] | |
| made_prediction_list = [] | |
| steric_clash_list = [] | |
| without_rec_overlap_list = [] | |
| start_time = time.time() | |
| successful_names_list = [] | |
| for i, name in enumerate(tqdm(names)): | |
| mol, _ = read_sdf_or_mol2(f"{args.renumbered_atoms_dir}/{name}.sdf", None) | |
| sm = Chem.MolToSmiles(mol) | |
| m_order = list(mol.GetPropsAsDict(includePrivate=True, includeComputed=True)['_smilesAtomOutputOrder']) | |
| mol = Chem.RenumberAtoms(mol, m_order) | |
| mol = Chem.RemoveHs(mol) | |
| orig_ligand_pos = np.array(mol.GetConformer().GetPositions()) | |
| assert(os.path.exists(os.path.join(args.results_path, name, f'{name}_tankbind_0.sdf'))) | |
| ligand_pos = [] | |
| for i in range(args.num_predictions): | |
| if not os.path.exists(os.path.join(args.results_path, name, f'{name}_tankbind_{i}.sdf')): break | |
| mol_pred, _ = read_sdf_or_mol2(os.path.join(args.results_path, name, f'{name}_tankbind_{i}.sdf'),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) | |
| ligand_pos.append(np.array(mol_pred.GetConformer().GetPositions())) | |
| ligand_pos = np.asarray(ligand_pos) | |
| try: | |
| unsym_rmsd = np.sqrt(((ligand_pos - orig_ligand_pos) ** 2).sum(axis=2).mean(axis=1)) | |
| rmsd = np.array(get_symmetry_rmsd(mol, orig_ligand_pos, [l for l in ligand_pos], mol_pred)) | |
| except Exception as e: | |
| print("Using non corrected RMSD because of the error:", e) | |
| rmsd = np.sqrt(((ligand_pos - orig_ligand_pos) ** 2).sum(axis=2).mean(axis=1)) | |
| num_pockets = len(ligand_pos) | |
| unsym_rmsds_list.append(np.lib.pad(unsym_rmsd, (0,10-len(unsym_rmsd)), 'constant', constant_values=(0)) ) | |
| rmsds_list.append(np.lib.pad(rmsd, (0,10-len(rmsd)), 'constant', constant_values=(0)) ) | |
| centroid_distance = np.linalg.norm(ligand_pos.mean(axis=1) - orig_ligand_pos[None,:].mean(axis=1), axis=1) | |
| centroid_distances_list.append(np.lib.pad(centroid_distance, (0,10-len(rmsd)), 'constant', constant_values=(0)) ) | |
| rec_path = os.path.join(args.data_dir, name, f'{name}_protein_processed.pdb') | |
| if not os.path.exists(rec_path): | |
| rec_path = os.path.join(args.data_dir, name,f'{name}_protein_obabel_reduce.pdb') | |
| rec = PandasPdb().read_pdb(rec_path) | |
| rec_df = rec.df['ATOM'] | |
| receptor_pos = rec_df[['x_coord', 'y_coord', 'z_coord']].to_numpy().squeeze().astype(np.float32) | |
| receptor_pos = np.tile(receptor_pos, (10, 1, 1)) | |
| ligand_pos_padded = np.lib.pad(ligand_pos, ((0,10-len(ligand_pos)), (0,0), (0,0)), 'constant', constant_values=(np.inf)) | |
| ligand_pos_padded_zero = np.lib.pad(ligand_pos, ((0, 10 - len(ligand_pos)), (0, 0), (0, 0)), 'constant',constant_values=0) | |
| cross_distances = np.linalg.norm(receptor_pos[:, :, None, :] - ligand_pos_padded[:, None, :, :], axis=-1) | |
| self_distances = np.linalg.norm(ligand_pos_padded_zero[:, :, None, :] - ligand_pos_padded_zero[:, None, :, :], axis=-1) | |
| self_distances = np.where(np.eye(self_distances.shape[2]), np.inf, self_distances) | |
| min_self_distances_list.append(np.min(self_distances, axis=(1, 2))) | |
| min_cross_distance = np.min(cross_distances, axis=(1, 2)) | |
| individual_made_prediction = np.lib.pad(np.ones(num_pockets), (0,10-len(rmsd)), 'constant', constant_values=(0)) | |
| made_prediction_list.append(individual_made_prediction) | |
| min_cross_distances_list.append(min_cross_distance) | |
| successful_names_list.append(name) | |
| without_rec_overlap_list.append(1 if name in names_no_rec_overlap else 0) | |
| performance_metrics = {} | |
| for overlap in ['', 'no_overlap_']: | |
| if 'no_overlap_' == overlap: | |
| without_rec_overlap = np.array(without_rec_overlap_list, dtype=bool) | |
| unsym_rmsds = np.array(unsym_rmsds_list)[without_rec_overlap] | |
| rmsds = np.array(rmsds_list)[without_rec_overlap] | |
| centroid_distances = np.array(centroid_distances_list)[without_rec_overlap] | |
| min_cross_distances = np.array(min_cross_distances_list)[without_rec_overlap] | |
| min_self_distances = np.array(min_self_distances_list)[without_rec_overlap] | |
| made_prediction = np.array(made_prediction_list)[without_rec_overlap] | |
| successful_names = np.array(successful_names_list)[without_rec_overlap] | |
| else: | |
| unsym_rmsds = np.array(unsym_rmsds_list) | |
| rmsds = np.array(rmsds_list) | |
| centroid_distances = np.array(centroid_distances_list) | |
| min_cross_distances = np.array(min_cross_distances_list) | |
| min_self_distances = np.array(min_self_distances_list) | |
| made_prediction = np.array(made_prediction_list) | |
| successful_names = np.array(successful_names_list) | |
| inf_rmsds = copy.deepcopy(rmsds) | |
| inf_rmsds[~made_prediction.astype(bool)] = np.inf | |
| inf_centroid_distances = copy.deepcopy(centroid_distances) | |
| inf_centroid_distances[~made_prediction.astype(bool)] = np.inf | |
| np.save(os.path.join(args.results_path, f'{overlap}rmsds.npy'), rmsds) | |
| np.save(os.path.join(args.results_path, f'{overlap}names.npy'), np.array(successful_names)) | |
| np.save(os.path.join(args.results_path, f'{overlap}centroid_distances.npy'), centroid_distances) | |
| np.save(os.path.join(args.results_path, f'{overlap}min_cross_distances.npy'), min_cross_distances) | |
| np.save(os.path.join(args.results_path, f'{overlap}min_self_distances.npy'), min_self_distances) | |
| performance_metrics.update({ | |
| f'{overlap}self_intersect_fraction': (100 * (min_self_distances[:, 0] < 0.4).sum() / len(min_self_distances[:, 0])), | |
| f'{overlap}steric_clash_fraction': (100 * (min_cross_distances[:,0] < 0.4).sum() / len(min_cross_distances[:,0])), | |
| f'{overlap}mean_rmsd': rmsds[:,0].mean(), | |
| f'{overlap}unsym_rmsds_below_2': (100 * (unsym_rmsds[:,0] < 2).sum() / len(unsym_rmsds[:,0])), | |
| f'{overlap}rmsds_below_2': (100 * (rmsds[:,0] < 2).sum() / len(rmsds[:,0])), | |
| f'{overlap}rmsds_below_5': (100 * (rmsds[:,0] < 5).sum() / len(rmsds[:,0])), | |
| f'{overlap}rmsds_percentile_25': np.percentile(rmsds[:,0], 25).round(2), | |
| f'{overlap}rmsds_percentile_50': np.percentile(rmsds[:,0], 50).round(2), | |
| f'{overlap}rmsds_percentile_75': np.percentile(rmsds[:,0], 75).round(2), | |
| f'{overlap}mean_centroid': centroid_distances[:,0].mean().__round__(2), | |
| f'{overlap}centroid_below_2': (100 * (centroid_distances[:,0] < 2).sum() / len(centroid_distances[:,0])).__round__(2), | |
| f'{overlap}centroid_below_5': (100 * (centroid_distances[:,0] < 5).sum() / len(centroid_distances[:,0])).__round__(2), | |
| f'{overlap}centroid_percentile_25': np.percentile(centroid_distances[:,0], 25).round(2), | |
| f'{overlap}centroid_percentile_50': np.percentile(centroid_distances[:,0], 50).round(2), | |
| f'{overlap}centroid_percentile_75': np.percentile(centroid_distances[:,0], 75).round(2), | |
| }) | |
| top5_rmsds = np.min(inf_rmsds[:, :5], axis=1) | |
| top5_centroid_distances = centroid_distances[np.arange(rmsds.shape[0])[:,None],np.argsort(inf_rmsds[:, :5], axis=1)][:,0] | |
| top5_min_cross_distances = min_cross_distances[np.arange(rmsds.shape[0])[:,None],np.argsort(inf_rmsds[:, :5], axis=1)][:,0] | |
| top5_min_self_distances = min_self_distances[np.arange(rmsds.shape[0])[:,None],np.argsort(inf_rmsds[:, :5], axis=1)][:,0] | |
| performance_metrics.update({ | |
| f'{overlap}top5_steric_clash_fraction': (100 * (top5_min_cross_distances < 0.4).sum() / len(top5_min_cross_distances)).__round__(2), | |
| f'{overlap}top5_self_intersect_fraction': (100 * (top5_min_self_distances < 0.4).sum() / len(top5_min_self_distances)).__round__(2), | |
| f'{overlap}top5_rmsds_below_2': (100 * (top5_rmsds < 2).sum() / len(top5_rmsds)).__round__(2), | |
| f'{overlap}top5_rmsds_below_5': (100 * (top5_rmsds < 5).sum() / len(top5_rmsds)).__round__(2), | |
| f'{overlap}top5_rmsds_percentile_25': np.percentile(top5_rmsds, 25).round(2), | |
| f'{overlap}top5_rmsds_percentile_50': np.percentile(top5_rmsds, 50).round(2), | |
| f'{overlap}top5_rmsds_percentile_75': np.percentile(top5_rmsds, 75).round(2), | |
| f'{overlap}top5_centroid_below_2': (100 * (top5_centroid_distances < 2).sum() / len(top5_centroid_distances)).__round__(2), | |
| f'{overlap}top5_centroid_below_5': (100 * (top5_centroid_distances < 5).sum() / len(top5_centroid_distances)).__round__(2), | |
| f'{overlap}top5_centroid_percentile_25': np.percentile(top5_centroid_distances, 25).round(2), | |
| f'{overlap}top5_centroid_percentile_50': np.percentile(top5_centroid_distances, 50).round(2), | |
| f'{overlap}top5_centroid_percentile_75': np.percentile(top5_centroid_distances, 75).round(2), | |
| }) | |
| top10_rmsds = np.min(inf_rmsds[:, :10], axis=1) | |
| top10_centroid_distances = centroid_distances[np.arange(rmsds.shape[0])[:,None],np.argsort(inf_rmsds[:, :10], axis=1)][:,0] | |
| top10_min_cross_distances = min_cross_distances[np.arange(rmsds.shape[0])[:,None],np.argsort(inf_rmsds[:, :10], axis=1)][:,0] | |
| top10_min_self_distances = min_self_distances[np.arange(rmsds.shape[0])[:,None],np.argsort(inf_rmsds[:, :10], axis=1)][:,0] | |
| performance_metrics.update({ | |
| f'{overlap}top10_steric_clash_fraction': (100 * (top10_min_cross_distances < 0.4).sum() / len(top10_min_cross_distances)).__round__(2), | |
| f'{overlap}top10_self_intersect_fraction': (100 * (top10_min_self_distances < 0.4).sum() / len(top10_min_self_distances)).__round__(2), | |
| f'{overlap}top10_rmsds_below_2': (100 * (top10_rmsds < 2).sum() / len(top10_rmsds)).__round__(2), | |
| f'{overlap}top10_rmsds_below_5': (100 * (top10_rmsds < 5).sum() / len(top10_rmsds)).__round__(2), | |
| f'{overlap}top10_rmsds_percentile_25': np.percentile(top10_rmsds, 25).round(2), | |
| f'{overlap}top10_rmsds_percentile_50': np.percentile(top10_rmsds, 50).round(2), | |
| f'{overlap}top10_rmsds_percentile_75': np.percentile(top10_rmsds, 75).round(2), | |
| f'{overlap}top10_centroid_below_2': (100 * (top10_centroid_distances < 2).sum() / len(top10_centroid_distances)).__round__(2), | |
| f'{overlap}top10_centroid_below_5': (100 * (top10_centroid_distances < 5).sum() / len(top10_centroid_distances)).__round__(2), | |
| f'{overlap}top10_centroid_percentile_25': np.percentile(top10_centroid_distances, 25).round(2), | |
| f'{overlap}top10_centroid_percentile_50': np.percentile(top10_centroid_distances, 50).round(2), | |
| f'{overlap}top10_centroid_percentile_75': np.percentile(top10_centroid_distances, 75).round(2), | |
| }) | |
| for k in performance_metrics: | |
| print(k, performance_metrics[k]) | |
| if args.wandb: | |
| wandb.log(performance_metrics) | |
| histogram_metrics_list = [('rmsd', rmsds[:,0]), | |
| ('centroid_distance', centroid_distances[:,0]), | |
| ('mean_rmsd', rmsds[:,0]), | |
| ('mean_centroid_distance', centroid_distances[:,0])] | |
| histogram_metrics_list.append(('top5_rmsds', top5_rmsds)) | |
| histogram_metrics_list.append(('top5_centroid_distances', top5_centroid_distances)) | |
| histogram_metrics_list.append(('top10_rmsds', top10_rmsds)) | |
| histogram_metrics_list.append(('top10_centroid_distances', top10_centroid_distances)) | |
| os.makedirs(f'.plotly_cache/baseline_cache', exist_ok=True) | |
| images = [] | |
| for metric_name, metric in histogram_metrics_list: | |
| d = {args.results_path: metric} | |
| df = pd.DataFrame(data=d) | |
| fig = px.ecdf(df, width=900, height=600, range_x=[0, 40]) | |
| fig.add_vline(x=2, annotation_text='2 A;', annotation_font_size=20, annotation_position="top right", | |
| line_dash='dash', line_color='firebrick', annotation_font_color='firebrick') | |
| fig.add_vline(x=5, annotation_text='5 A;', annotation_font_size=20, annotation_position="top right", | |
| line_dash='dash', line_color='green', annotation_font_color='green') | |
| fig.update_xaxes(title=f'{metric_name} in Angstrom', title_font={"size": 20}, tickfont={"size": 20}) | |
| fig.update_yaxes(title=f'Fraction of predictions with lower error', title_font={"size": 20}, | |
| tickfont={"size": 20}) | |
| fig.update_layout(autosize=False, margin={'l': 0, 'r': 0, 't': 0, 'b': 0}, plot_bgcolor='white', | |
| paper_bgcolor='white', legend_title_text='Method', legend_title_font_size=17, | |
| legend=dict(yanchor="bottom", y=0.1, xanchor="right", x=0.99, font=dict(size=17), ), ) | |
| fig.update_xaxes(showgrid=True, gridcolor='lightgrey') | |
| fig.update_yaxes(showgrid=True, gridcolor='lightgrey') | |
| fig.write_image(os.path.join(f'.plotly_cache/baseline_cache', f'{metric_name}.png')) | |
| wandb.log({metric_name: wandb.Image(os.path.join(f'.plotly_cache/baseline_cache', f'{metric_name}.png'), caption=f"{metric_name}")}) | |
| images.append(wandb.Image(os.path.join(f'.plotly_cache/baseline_cache', f'{metric_name}.png'), caption=f"{metric_name}")) | |
| wandb.log({'images': images}) |