File size: 9,636 Bytes
3ac1d94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | import argparse
import os
import numpy as np
from rdkit import Chem
from rdkit import RDLogger
import torch
from tqdm.auto import tqdm
from glob import glob
from collections import Counter
from onescience.utils.targetdiff.evaluation import eval_atom_type, scoring_func, analyze, eval_bond_length
from onescience.utils.targetdiff import misc, reconstruct, transforms
from onescience.utils.targetdiff.evaluation.docking_qvina import QVinaDockingTask
from onescience.utils.targetdiff.evaluation.docking_vina import VinaDockingTask
def print_dict(d, logger):
for k, v in d.items():
if v is not None:
logger.info(f'{k}:\t{v:.4f}')
else:
logger.info(f'{k}:\tNone')
def print_ring_ratio(all_ring_sizes, logger):
for ring_size in range(3, 10):
n_mol = 0
for counter in all_ring_sizes:
if ring_size in counter:
n_mol += 1
logger.info(f'ring size: {ring_size} ratio: {n_mol / len(all_ring_sizes):.3f}')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('sample_path', type=str)
parser.add_argument('--verbose', type=eval, default=False)
parser.add_argument('--eval_step', type=int, default=-1)
parser.add_argument('--eval_num_examples', type=int, default=None)
parser.add_argument('--save', type=eval, default=True)
parser.add_argument('--protein_root', type=str, default='./data/crossdocked_v1.1_rmsd1.0')
parser.add_argument('--atom_enc_mode', type=str, default='add_aromatic')
parser.add_argument('--docking_mode', type=str, choices=['qvina', 'vina_score', 'vina_dock', 'none'])
parser.add_argument('--exhaustiveness', type=int, default=16)
args = parser.parse_args()
result_path = os.path.join(args.sample_path, 'eval_results')
os.makedirs(result_path, exist_ok=True)
logger = misc.get_logger('evaluate', log_dir=result_path)
if not args.verbose:
RDLogger.DisableLog('rdApp.*')
# Load generated data
results_fn_list = glob(os.path.join(args.sample_path, '*result_*.pt'))
results_fn_list = sorted(results_fn_list, key=lambda x: int(os.path.basename(x)[:-3].split('_')[-1]))
if args.eval_num_examples is not None:
results_fn_list = results_fn_list[:args.eval_num_examples]
num_examples = len(results_fn_list)
logger.info(f'Load generated data done! {num_examples} examples in total.')
num_samples = 0
all_mol_stable, all_atom_stable, all_n_atom = 0, 0, 0
n_recon_success, n_eval_success, n_complete = 0, 0, 0
results = []
all_pair_dist, all_bond_dist = [], []
all_atom_types = Counter()
success_pair_dist, success_atom_types = [], Counter()
for example_idx, r_name in enumerate(tqdm(results_fn_list, desc='Eval')):
r = torch.load(r_name) # ['data', 'pred_ligand_pos', 'pred_ligand_v', 'pred_ligand_pos_traj', 'pred_ligand_v_traj']
all_pred_ligand_pos = r['pred_ligand_pos_traj'] # [num_samples, num_steps, num_atoms, 3]
all_pred_ligand_v = r['pred_ligand_v_traj']
num_samples += len(all_pred_ligand_pos)
for sample_idx, (pred_pos, pred_v) in enumerate(zip(all_pred_ligand_pos, all_pred_ligand_v)):
pred_pos, pred_v = pred_pos[args.eval_step], pred_v[args.eval_step]
# stability check
pred_atom_type = transforms.get_atomic_number_from_index(pred_v, mode=args.atom_enc_mode)
all_atom_types += Counter(pred_atom_type)
r_stable = analyze.check_stability(pred_pos, pred_atom_type)
all_mol_stable += r_stable[0]
all_atom_stable += r_stable[1]
all_n_atom += r_stable[2]
pair_dist = eval_bond_length.pair_distance_from_pos_v(pred_pos, pred_atom_type)
all_pair_dist += pair_dist
# reconstruction
try:
pred_aromatic = transforms.is_aromatic_from_index(pred_v, mode=args.atom_enc_mode)
mol = reconstruct.reconstruct_from_generated(pred_pos, pred_atom_type, pred_aromatic)
smiles = Chem.MolToSmiles(mol)
except reconstruct.MolReconsError:
if args.verbose:
logger.warning('Reconstruct failed %s' % f'{example_idx}_{sample_idx}')
continue
n_recon_success += 1
if '.' in smiles:
continue
n_complete += 1
# chemical and docking check
try:
chem_results = scoring_func.get_chem(mol)
if args.docking_mode == 'qvina':
vina_task = QVinaDockingTask.from_generated_mol(
mol, r['data'].ligand_filename, protein_root=args.protein_root)
vina_results = vina_task.run_sync()
elif args.docking_mode in ['vina_score', 'vina_dock']:
vina_task = VinaDockingTask.from_generated_mol(
mol, r['data'].ligand_filename, protein_root=args.protein_root)
score_only_results = vina_task.run(mode='score_only', exhaustiveness=args.exhaustiveness)
minimize_results = vina_task.run(mode='minimize', exhaustiveness=args.exhaustiveness)
vina_results = {
'score_only': score_only_results,
'minimize': minimize_results
}
if args.docking_mode == 'vina_dock':
docking_results = vina_task.run(mode='dock', exhaustiveness=args.exhaustiveness)
vina_results['dock'] = docking_results
else:
vina_results = None
n_eval_success += 1
except:
if args.verbose:
logger.warning('Evaluation failed for %s' % f'{example_idx}_{sample_idx}')
continue
# now we only consider complete molecules as success
bond_dist = eval_bond_length.bond_distance_from_mol(mol)
all_bond_dist += bond_dist
success_pair_dist += pair_dist
success_atom_types += Counter(pred_atom_type)
results.append({
'mol': mol,
'smiles': smiles,
'ligand_filename': r['data'].ligand_filename,
'pred_pos': pred_pos,
'pred_v': pred_v,
'chem_results': chem_results,
'vina': vina_results
})
logger.info(f'Evaluate done! {num_samples} samples in total.')
fraction_mol_stable = all_mol_stable / num_samples
fraction_atm_stable = all_atom_stable / all_n_atom
fraction_recon = n_recon_success / num_samples
fraction_eval = n_eval_success / num_samples
fraction_complete = n_complete / num_samples
validity_dict = {
'mol_stable': fraction_mol_stable,
'atm_stable': fraction_atm_stable,
'recon_success': fraction_recon,
'eval_success': fraction_eval,
'complete': fraction_complete
}
print_dict(validity_dict, logger)
c_bond_length_profile = eval_bond_length.get_bond_length_profile(all_bond_dist)
c_bond_length_dict = eval_bond_length.eval_bond_length_profile(c_bond_length_profile)
logger.info('JS bond distances of complete mols: ')
print_dict(c_bond_length_dict, logger)
success_pair_length_profile = eval_bond_length.get_pair_length_profile(success_pair_dist)
success_js_metrics = eval_bond_length.eval_pair_length_profile(success_pair_length_profile)
print_dict(success_js_metrics, logger)
atom_type_js = eval_atom_type.eval_atom_type_distribution(success_atom_types)
logger.info('Atom type JS: %.4f' % atom_type_js)
if args.save:
eval_bond_length.plot_distance_hist(success_pair_length_profile,
metrics=success_js_metrics,
save_path=os.path.join(result_path, f'pair_dist_hist_{args.eval_step}.png'))
logger.info('Number of reconstructed mols: %d, complete mols: %d, evaluated mols: %d' % (
n_recon_success, n_complete, len(results)))
qed = [r['chem_results']['qed'] for r in results]
sa = [r['chem_results']['sa'] for r in results]
logger.info('QED: Mean: %.3f Median: %.3f' % (np.mean(qed), np.median(qed)))
logger.info('SA: Mean: %.3f Median: %.3f' % (np.mean(sa), np.median(sa)))
if args.docking_mode == 'qvina':
vina = [r['vina'][0]['affinity'] for r in results]
logger.info('Vina: Mean: %.3f Median: %.3f' % (np.mean(vina), np.median(vina)))
elif args.docking_mode in ['vina_dock', 'vina_score']:
vina_score_only = [r['vina']['score_only'][0]['affinity'] for r in results]
vina_min = [r['vina']['minimize'][0]['affinity'] for r in results]
logger.info('Vina Score: Mean: %.3f Median: %.3f' % (np.mean(vina_score_only), np.median(vina_score_only)))
logger.info('Vina Min : Mean: %.3f Median: %.3f' % (np.mean(vina_min), np.median(vina_min)))
if args.docking_mode == 'vina_dock':
vina_dock = [r['vina']['dock'][0]['affinity'] for r in results]
logger.info('Vina Dock : Mean: %.3f Median: %.3f' % (np.mean(vina_dock), np.median(vina_dock)))
# check ring distribution
print_ring_ratio([r['chem_results']['ring_size'] for r in results], logger)
if args.save:
torch.save({
'stability': validity_dict,
'bond_length': all_bond_dist,
'all_results': results
}, os.path.join(result_path, f'metrics_{args.eval_step}.pt'))
|