| import argparse |
| import os |
| import torch |
| from tqdm.auto import tqdm |
| from onescience.utils.targetdiff.evaluation.docking_qvina import QVinaDockingTask |
| from onescience.utils.targetdiff.evaluation.docking_vina import VinaDockingTask |
| import multiprocessing as mp |
|
|
|
|
| def dock_pocket_samples(pocket_samples): |
| ligand_fn = pocket_samples[0]['ligand_filename'] |
| print('Start docking pocket: %s' % ligand_fn) |
| pocket_results = [] |
| for idx, s in enumerate(tqdm(pocket_samples, desc='docking %d' % os.getpid())): |
| try: |
| if args.docking_mode == 'qvina': |
| vina_task = QVinaDockingTask.from_generated_mol( |
| s['mol'], s['ligand_filename'], protein_root=args.protein_root, size_factor=args.dock_size_factor) |
| vina_results = vina_task.run_sync() |
| elif args.docking_mode == 'vina_score': |
| vina_task = VinaDockingTask.from_generated_mol( |
| s['mol'], s['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 |
| } |
| else: |
| raise ValueError |
| except: |
| print('Error at %d of %s' % (idx, ligand_fn)) |
| vina_results = None |
| pocket_results.append({**s, 'vina': vina_results}) |
| return pocket_results |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('sample_path', type=str) |
| parser.add_argument('-o', '--out', type=str, default=None) |
| parser.add_argument('-n', '--num_processes', type=int, default=10) |
| parser.add_argument('--protein_root', type=str, default='./data/crossdocked_v1.1_rmsd1.0') |
| parser.add_argument('--dock_size_factor', type=float, default=None) |
| parser.add_argument('--exhaustiveness', type=int, default=16) |
| parser.add_argument('--docking_mode', type=str, default='vina_score', |
| choices=['none', 'qvina', 'vina_score']) |
| args = parser.parse_args() |
|
|
| samples = torch.load(args.sample_path) |
| with mp.Pool(args.num_processes) as p: |
| docked_samples = p.map(dock_pocket_samples, samples) |
| if args.out is None: |
| dir_name = os.path.dirname(args.sample_path) |
| baseline_name = os.path.basename(args.sample_path).split('_')[0] |
| out_path = os.path.join(dir_name, baseline_name + '_test_docked.pt') |
| else: |
| out_path = args.out |
| torch.save(docked_samples, out_path) |
|
|
