import os import re import tempfile import numpy as np import torch from pathlib import Path import argparse import pandas as pd from rdkit import Chem from tqdm import tqdm try: import utils except ModuleNotFoundError as e: print(e) def calculate_smina_score(pdb_file, sdf_file): # add '-o _smina.sdf' if you want to see the output out = os.popen(f'smina.static -l {sdf_file} -r {pdb_file} ' f'--score_only').read() matches = re.findall( r"Affinity:[ ]+([+-]?[0-9]*[.]?[0-9]+)[ ]+\(kcal/mol\)", out) return [float(x) for x in matches] def smina_score(rdmols, receptor_file): """ Calculate smina score :param rdmols: List of RDKit molecules :param receptor_file: Receptor pdb/pdbqt file or list of receptor files :return: Smina score for each input molecule (list) """ if isinstance(receptor_file, list): scores = [] for mol, rec_file in zip(rdmols, receptor_file): with tempfile.NamedTemporaryFile(suffix='.sdf') as tmp: tmp_file = tmp.name utils.write_sdf_file(tmp_file, [mol]) scores.extend(calculate_smina_score(rec_file, tmp_file)) # Use same receptor file for all molecules else: with tempfile.NamedTemporaryFile(suffix='.sdf') as tmp: tmp_file = tmp.name utils.write_sdf_file(tmp_file, rdmols) scores = calculate_smina_score(receptor_file, tmp_file) return scores def sdf_to_pdbqt(sdf_file, pdbqt_outfile, mol_id): os.popen(f'obabel {sdf_file} -O {pdbqt_outfile} ' f'-f {mol_id + 1} -l {mol_id + 1}').read() return pdbqt_outfile def calculate_qvina2_score(receptor_file, sdf_file, out_dir, size=20, exhaustiveness=16, return_rdmol=False): receptor_file = Path(receptor_file) sdf_file = Path(sdf_file) if receptor_file.suffix == '.pdb': # prepare receptor, requires Python 2.7 receptor_pdbqt_file = Path(out_dir, receptor_file.stem + '.pdbqt') os.popen(f'prepare_receptor4.py -r {receptor_file} -O {receptor_pdbqt_file}') else: receptor_pdbqt_file = receptor_file scores = [] rdmols = [] # for if return rdmols suppl = Chem.SDMolSupplier(str(sdf_file), sanitize=False) for i, mol in enumerate(suppl): # sdf file may contain several ligands ligand_name = f'{sdf_file.stem}_{i}' # prepare ligand ligand_pdbqt_file = Path(out_dir, ligand_name + '.pdbqt') out_sdf_file = Path(out_dir, ligand_name + '_out.sdf') if out_sdf_file.exists(): with open(out_sdf_file, 'r') as f: scores.append( min([float(x.split()[2]) for x in f.readlines() if x.startswith(' VINA RESULT:')]) ) else: sdf_to_pdbqt(sdf_file, ligand_pdbqt_file, i) # center box at ligand's center of mass cx, cy, cz = mol.GetConformer().GetPositions().mean(0) # run QuickVina 2 out = os.popen( f'qvina2.1 --receptor {receptor_pdbqt_file} ' f'--ligand {ligand_pdbqt_file} ' f'--center_x {cx:.4f} --center_y {cy:.4f} --center_z {cz:.4f} ' f'--size_x {size} --size_y {size} --size_z {size} ' f'--exhaustiveness {exhaustiveness}' ).read() # clean up ligand_pdbqt_file.unlink() if '-----+------------+----------+----------' not in out: scores.append(np.nan) continue out_split = out.splitlines() best_idx = out_split.index('-----+------------+----------+----------') + 1 best_line = out_split[best_idx].split() assert best_line[0] == '1' scores.append(float(best_line[1])) out_pdbqt_file = Path(out_dir, ligand_name + '_out.pdbqt') if out_pdbqt_file.exists(): os.popen(f'obabel {out_pdbqt_file} -O {out_sdf_file}').read() # clean up out_pdbqt_file.unlink() if return_rdmol: rdmol = Chem.SDMolSupplier(str(out_sdf_file))[0] rdmols.append(rdmol) if return_rdmol: return scores, rdmols else: return scores if __name__ == '__main__': parser = argparse.ArgumentParser('QuickVina evaluation') parser.add_argument('--pdbqt_dir', type=Path, help='Receptor files in pdbqt format') parser.add_argument('--sdf_dir', type=Path, default=None, help='Ligand files in sdf format') parser.add_argument('--sdf_files', type=Path, nargs='+', default=None) parser.add_argument('--out_dir', type=Path) parser.add_argument('--write_csv', action='store_true') parser.add_argument('--write_dict', action='store_true') parser.add_argument('--dataset', type=str, default='moad') args = parser.parse_args() assert (args.sdf_dir is not None) ^ (args.sdf_files is not None) args.out_dir.mkdir(exist_ok=True) results = {'receptor': [], 'ligand': [], 'scores': []} results_dict = {} sdf_files = list(args.sdf_dir.glob('[!.]*.sdf')) \ if args.sdf_dir is not None else args.sdf_files pbar = tqdm(sdf_files) for sdf_file in pbar: pbar.set_description(f'Processing {sdf_file.name}') if args.dataset == 'moad': """ Ligand file names should be of the following form: __.sdf where and cannot contain any underscores, e.g.: 1abc-bio1_pocket0_gen.sdf """ ligand_name = sdf_file.stem receptor_name, pocket_id, *suffix = ligand_name.split('_') suffix = '_'.join(suffix) receptor_file = Path(args.pdbqt_dir, receptor_name + '.pdbqt') elif args.dataset == 'crossdocked': ligand_name = sdf_file.stem receptor_name = ligand_name[:-4] receptor_file = Path(args.pdbqt_dir, receptor_name + '.pdbqt') # try: scores, rdmols = calculate_qvina2_score( receptor_file, sdf_file, args.out_dir, return_rdmol=True) # except AttributeError as e: # print(e) # continue results['receptor'].append(str(receptor_file)) results['ligand'].append(str(sdf_file)) results['scores'].append(scores) if args.write_dict: results_dict[ligand_name] = { 'receptor': str(receptor_file), 'ligand': str(sdf_file), 'scores': scores, 'rmdols': rdmols } if args.write_csv: df = pd.DataFrame.from_dict(results) df.to_csv(Path(args.out_dir, 'qvina2_scores.csv')) if args.write_dict: torch.save(results_dict, Path(args.out_dir, 'qvina2_scores.pt'))