| import pickle |
| import os |
| import glob |
| from multiprocessing import Pool |
| import numpy as np |
| from rdkit import Chem |
| from scipy.spatial import distance_matrix |
| from Bio.PDB import * |
| from Bio.PDB.PDBIO import Select |
| import warnings |
| warnings.filterwarnings('ignore') |
|
|
| def extract(ligand, pdb,key): |
| parser = PDBParser() |
| structure = parser.get_structure("protein", pdb) |
| ligand_positions = ligand.GetConformer().GetPositions() |
| |
| |
| |
| class ResidueSelect(Select): |
| def accept_residue(self, residue): |
| residue_positions = np.array([np.array(list(atom.get_vector())) \ |
| for atom in residue.get_atoms()]) |
| if len(residue_positions.shape) < 2: |
| print(residue) |
| return 0 |
| min_dis = np.min(distance_matrix(residue_positions, ligand_positions)) |
| if min_dis < 8.0: |
| return 1 |
| else: |
| return 0 |
| |
| io = PDBIO() |
| io.set_structure(structure) |
| fn = "BS_tmp_"+str(key)+".pdb" |
| io.save(fn, ResidueSelect()) |
| try: |
| m2 = Chem.MolFromPDBFile(fn) |
| |
| if m2 is None: |
| print("first read PDB fail",fn) |
| |
| remove_zn_dir="./docker_result_remove_ZN" |
| if not os.path.exists(remove_zn_dir): |
| os.mkdir(remove_zn_dir) |
| cmd=f"cp {fn} {remove_zn_dir}" |
| print(cmd) |
| os.system(cmd) |
| fn_remove_zn=os.path.join(remove_zn_dir,fn.replace('.pdb','_remove_ZN.pdb')) |
| cmd=f"sed -e '/ZN/d' {fn} > {fn_remove_zn}" |
| os.system(cmd) |
| print("delete metal atom and get new pdb file",fn_remove_zn) |
| m2 = Chem.MolFromPDBFile(fn_remove_zn) |
| else: |
| os.system("rm -f " + fn) |
| except: |
| print("Read PDB fail for other unknow reason",fn) |
| return m2 |
|
|
| def preprocessor(docking_result_sdf_fn,origin_recptor_pdb,data_dir): |
| """ |
| get pocket from docking result and save to file:(m1,m2) |
| |
| input: |
| docking_result_sdf_fn: docking result sdf file, one ligand in sdf file will speed up this process in multi-process |
| origin_recptor_pdb: receptor pdb file |
| data_dir: path for save pocket file |
| output: |
| 0: success |
| -1: fail |
| """ |
| sdf_fn = docking_result_sdf_fn.split("/")[-1].split(".")[0] |
| if not os.path.exists(data_dir): |
| os.mkdir(data_dir) |
| if os.path.getsize(docking_result_sdf_fn): |
| total=Chem.SDMolSupplier(docking_result_sdf_fn) |
| for i,m1 in enumerate(total): |
| key=sdf_fn |
| if not os.path.exists(os.path.join(data_dir,key)): |
|
|
| if len(m1.GetConformers())==0: |
| print(f"{key} mol no conformer!") |
| continue |
| try: |
| m2 = extract(m1, origin_recptor_pdb,key) |
| except: |
| print(f'extract m2 failed {sdf_fn}') |
| continue |
|
|
| if m2 is None : |
| print(f"{key} no extracted binding pocket!") |
| continue |
| if len(m2.GetConformers())==0: |
| print(f"{key} receptor no conformer!") |
| continue |
|
|
| with open(os.path.join(data_dir,key), "wb") as fp: |
| pickle.dump((m1, m2), fp, pickle.HIGHEST_PROTOCOL) |
| else: |
| print(f'file done before so skip it {sdf_fn}') |
| continue |
| return 0 |
| |
| else: |
| print("docking result file size is 0!") |
| return -1 |
| def out_sdf(lig,filename): |
| writer = Chem.SDWriter(filename) |
| writer.write(lig) |
| writer.close() |
| return |
| def get_pocket_with_water(complex_sample,receptor_fn): |
| status=preprocessor(complex_sample,receptor_fn,out_data_dir) |
| |
| if __name__ == '__main__': |
|
|
| import time |
| from multiprocessing import Pool |
| import os |
| import gzip |
| import tqdm |
| |
| import argparse |
| parser = argparse.ArgumentParser(description='Process data from docking result') |
| parser.add_argument("--single_sdf_save_path", help="file path for save compounds from docking result.", type=str, \ |
| default=None,required=True) |
| parser.add_argument("--docking_result", help="docking result filname.maegz,filename.mae or filename.sdf.", type=str,default=None,required=True) |
| parser.add_argument("--recptor_pdb", help="receptor pdb file.", type=str,default=None,required=True) |
| parser.add_argument("--pocket_save_dir", help="save pocket file dir.", type=str,default=None,required=True) |
| parser.add_argument("--prefix", help="Anything that helps you distinguish between compounds.", type=str,default='Compound') |
| parser.add_argument("--process_num", help="process num for multi process ", type=int,default=1) |
| args = parser.parse_args() |
| os.makedirs(args.single_sdf_save_path,exist_ok=True) |
| if args.docking_result.endswith('maegz'): |
| total=Chem.rdmolfiles.MaeMolSupplier(gzip.open(args.docking_result)) |
| elif args.docking_result.endswith('sdf'): |
| total=Chem.SDMolSupplier(args.docking_result) |
| elif args.docking_result.endswith('mae'): |
| total=Chem.rdmolfiles.MaeMolSupplier(args.docking_result) |
| else: |
| print('docking result file format error! only support maegz,mae or sdf') |
| exit() |
| |
| for i,sample in enumerate(total): |
| if i==0 and len(sample.GetAtoms()) > 500: |
| print('atoms nums',len(sample.GetAtoms()),'may you not split protein and compounds ? save protein in a file in this dir') |
| Chem.MolToPDBFile(sample,f'./data/protein.pdb') |
| print('save protein success') |
| else: |
| if sample is not None: |
| name = '{}_{}_{}.sdf'.format(os.path.basename(args.docking_result).split('.')[0],args.prefix,i) |
| out_sdf(sample,os.path.join(args.single_sdf_save_path,name)) |
| |
| total_sdfs = [os.path.join(args.single_sdf_save_path,filename)for filename in os.listdir(args.single_sdf_save_path)] |
| file_tuple_list = [] |
| for complex_sample in total_sdfs: |
| receptor_fn=args.recptor_pdb |
| file_tuple_list.append((complex_sample,receptor_fn)) |
| print('num compounds to get pocket',len(file_tuple_list)) |
| out_data_dir = args.pocket_save_dir |
| p = Pool(args.process_num) |
| pbar = tqdm.tqdm(total=len(file_tuple_list)) |
| pbar.set_description('get_pocket:') |
| update = lambda *args: pbar.update() |
| for file_tuple in file_tuple_list: |
| p.apply_async(get_pocket_with_water,args = (file_tuple[0],file_tuple[1]),callback=update) |
| print('waiting for processing!') |
| p.close() |
| p.join() |
| print("all pocket done! check the outdir plz!") |
