baselines/baseline_tankbind_evaluation.py

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