Spaces:
Sleeping
Sleeping
| import os | |
| import subprocess | |
| import numpy as np | |
| from rdkit.Chem import AllChem, RemoveHs, RemoveAllHs | |
| from datasets.process_mols import write_mol_with_coords, read_molecule | |
| import re | |
| from utils.utils import remove_all_hs | |
| def read_gnina_metrics(gnina_sdf_path): | |
| with open(gnina_sdf_path, 'r') as f: | |
| pattern = re.compile(r'> <(.*?)>\n(.*?)\n') | |
| content = f.read() | |
| matches = pattern.findall(content) | |
| metrics = {k: v for k, v in matches} | |
| return metrics | |
| def read_gnina_score(gnina_sdf_path): | |
| with open(gnina_sdf_path, 'r') as f: | |
| pattern = re.compile(r'> <CNNscore>\n(.*?)\n') | |
| content = f.read() | |
| matches = pattern.findall(content) | |
| return float(matches[0]) | |
| def invert_permutation(p): | |
| """Return an array s with which np.array_equal(arr[p][s], arr) is True. | |
| The array_like argument p must be some permutation of 0, 1, ..., len(p)-1. | |
| """ | |
| p = np.asanyarray(p) # in case p is a tuple, etc. | |
| s = np.empty_like(p) | |
| s[p] = np.arange(p.size) | |
| return s | |
| def get_gnina_poses(args, mol, pos, orig_center, name, folder, gnina_path, thread_id=0): | |
| #folder = "data/MOAD_new_test_processed" if args.split == 'test' else "data/MOAD_new_val_processed" | |
| out_dir = args.out_dir if hasattr(args, 'out_dir') else args.inference_out_dir | |
| rec_path = os.path.join(folder, name[:6] + '_protein_chain_removed.pdb') | |
| pred_lig_path = os.path.join(out_dir, f'pred_{name}_tid{thread_id}_lig.sdf') | |
| if not os.path.exists(os.path.dirname(pred_lig_path)): | |
| os.mkdir(os.path.dirname(pred_lig_path)) | |
| print(f'Ligand path {pred_lig_path}') | |
| write_mol_with_coords(mol, pos + orig_center, pred_lig_path) | |
| gnina_pred_path = os.path.join(out_dir, f'gnina_{name}_tid{thread_id}_lig.sdf') | |
| gnina_logs_dir = os.path.join(out_dir, "gnina_logs") | |
| with open(os.path.join(gnina_logs_dir, f'{name}'), "w+") as f: | |
| if args.gnina_full_dock: | |
| return_code = subprocess.run( | |
| f'{gnina_path} -r {rec_path} -l "{pred_lig_path}" --autobox_ligand "{pred_lig_path}" -o "{gnina_pred_path}" --no_gpu --autobox_add {args.gnina_autobox_add}', | |
| shell=True, stdout=f, stderr=f) | |
| else: | |
| return_code = subprocess.run( | |
| f'{gnina_path} --receptor {rec_path} --ligand "{pred_lig_path}" --minimize -o "{gnina_pred_path}"', | |
| shell=True, stdout=f, stderr=f) | |
| # print(f'gnina return code: {return_code}') | |
| try: | |
| gnina_mol = RemoveAllHs(read_molecule(gnina_pred_path, remove_hs=True, sanitize=True)) | |
| gnina_minimized_ligand_pos = np.array(gnina_mol.GetConformer(0).GetPositions()) | |
| gnina_atoms = np.array([atom.GetSymbol() for atom in gnina_mol.GetAtoms()]) | |
| gnina_filter_Hs = np.where(gnina_atoms != 'H') | |
| gnina_ligand_pos = gnina_minimized_ligand_pos[gnina_filter_Hs] - orig_center | |
| try: | |
| gnina_score = read_gnina_score(gnina_pred_path) | |
| if gnina_score is None: | |
| gnina_score = 0 | |
| except Exception as e: | |
| print(f'Error reading gnina score: {e}') | |
| gnina_score = 0 | |
| except Exception as e: | |
| print(f'Error when running gnina with {name} to minimize energy') | |
| print('Error:', e) | |
| print('Using score model output pos instead.') | |
| gnina_ligand_pos = pos | |
| gnina_mol = RemoveAllHs(mol) | |
| gnina_score = 0 | |
| return gnina_ligand_pos, gnina_mol, gnina_score | |