File size: 7,132 Bytes
96f168d | 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 | 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()
# Get distance between ligand positions (N_ligand, 3) and
# residue positions (N_residue, 3) for each residue
# only select residue with minimum distance of it is smaller than 8A
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 "H" not in atom.get_id()
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)
# may contain metal atom, causing MolFromPDBFile return None
if m2 is None:
print("first read PDB fail",fn)
# copy file to tmp dir
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): #docking ligand file may be 0 size
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)
# print(status)
if __name__ == '__main__':
import time
from multiprocessing import Pool
import os
import gzip
import tqdm
# get pocket and save to file
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() # set callback function to update pbar state when process end
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!") |