File size: 8,233 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 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | 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')
from rdkit.Chem import AllChem
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(ligand_dir,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
"""
file_flag = os.path.basename(ligand_dir)
try:
m1 = read_molecule(os.path.join(ligand_dir, f'{file_flag}_ligand.sdf'), remove_hs=True, sanitize=True)
if m1 is None: # read mol2 file if sdf file cannot be sanitized
print('Using the .sdf file failed. We found a .mol2 file instead and are trying to use that.')
m1 = read_molecule(os.path.join(ligand_dir, f'{file_flag}_ligand.mol2'), remove_hs=True, sanitize=True)
except Exception as e:
print(e)
return -1
if not os.path.exists(data_dir):
os.mkdir(data_dir)
if m1 is not None: #docking ligand file may be 0 size
if not os.path.exists(os.path.join(data_dir,file_flag)):
os.mkdir(os.path.join(data_dir,file_flag))
if len(m1.GetConformers())==0:
print(f"{file_flag} mol no conformer!")
return -1
try:
pdb_path = os.path.join(ligand_dir, f'{file_flag}_protein_processed.pdb')
m2 = extract(m1,pdb_path ,file_flag)
except:
print(f'extract m2 failed {file_flag}')
return -1
if m2 is None :
print(f"{file_flag} no extracted binding pocket!")
# continue
return -1
if len(m2.GetConformers())==0:
print(f"{file_flag} receptor no conformer!")
return -1
# save pdb pocket
Chem.MolToPDBFile(m2, os.path.join(data_dir,file_flag,f'{file_flag}_pocket.pdb'))
else:
print(f'file done before so skip it {file_flag}')
return 0
# return 0
else:
print("read mol fail")
return -1
def out_sdf(lig,filename):
writer = Chem.SDWriter(filename)
writer.write(lig)
writer.close()
return
def get_pocket_with_water(complex_sample):
status=preprocessor(complex_sample,out_data_dir)
# print(status)
def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=False):
if molecule_file.endswith('.mol2'):
mol = Chem.MolFromMol2File(molecule_file, sanitize=False, removeHs=False)
elif molecule_file.endswith('.sdf'):
supplier = Chem.SDMolSupplier(molecule_file, sanitize=False, removeHs=False)
mol = supplier[0]
elif molecule_file.endswith('.pdbqt'):
with open(molecule_file) as file:
pdbqt_data = file.readlines()
pdb_block = ''
for line in pdbqt_data:
pdb_block += '{}\n'.format(line[:66])
mol = Chem.MolFromPDBBlock(pdb_block, sanitize=False, removeHs=False)
elif molecule_file.endswith('.pdb'):
mol = Chem.MolFromPDBFile(molecule_file, sanitize=False, removeHs=False)
else:
raise ValueError('Expect the format of the molecule_file to be '
'one of .mol2, .sdf, .pdbqt and .pdb, got {}'.format(molecule_file))
try:
if sanitize or calc_charges:
Chem.SanitizeMol(mol)
if calc_charges:
# Compute Gasteiger charges on the molecule.
try:
AllChem.ComputeGasteigerCharges(mol)
except:
warnings.warn('Unable to compute charges for the molecule.')
if remove_hs:
mol = Chem.RemoveHs(mol, sanitize=sanitize)
except Exception as e:
print(e)
print("RDKit was unable to read the molecule.")
return None
return mol
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("--PDBbind_path", help="file path for save compounds from docking result.", type=str, \
default='/home/house/caoduanhua/DeepLearningForDock/datasets/equibind_and_diffdock_dataset/PDBBIND/PDBBind_processed/',required=False)
# 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("--save_dir", help="save pocket file dir.", type=str,default='/home/house/caoduanhua/DeepLearningForDock/datasets/equibind_and_diffdock_dataset/PDBBIND/PDBBind_pocket_8A',required=False)
# parser.add_argument("--prefix", help="Anything that helps you distinguish between compounds.", type=str,default='')
parser.add_argument("--process_num", help="process num for multi process ", type=int,default=60)
args = parser.parse_args()
total_sdfs = [os.path.join(args.PDBbind_path,filename) for filename in os.listdir(args.PDBbind_path)]
file_tuple_list = []
for complex_sample in total_sdfs:
# receptor_fn=args.recptor_pdb
file_tuple_list.append(complex_sample)
print('num compounds to get pocket',len(file_tuple_list))
out_data_dir = args.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,),callback=update)
print('waiting for processing!')
p.close()
p.join()
print("all pocket done! check the outdir plz!") |