OneScience-Group/SurfDock

1import pickle2import os3import glob4from multiprocessing import Pool5import numpy as np6from rdkit import Chem7from scipy.spatial import distance_matrix8from Bio.PDB import *9from Bio.PDB.PDBIO import Select10import warnings11warnings.filterwarnings('ignore')12from rdkit.Chem import AllChem13def extract(ligand, pdb,key):14    parser = PDBParser()15    structure = parser.get_structure("protein", pdb)16    ligand_positions = ligand.GetConformer().GetPositions()17    # Get distance between ligand positions (N_ligand, 3) and18    # residue positions (N_residue, 3) for each residue19    # only select residue with minimum distance of it is smaller than 8A20    class ResidueSelect(Select):21        def accept_residue(self, residue):22            residue_positions = np.array([np.array(list(atom.get_vector())) \23                for atom in residue.get_atoms()])    # if "H" not in atom.get_id()24            if len(residue_positions.shape) < 2:25                print(residue)26                return 027            min_dis = np.min(distance_matrix(residue_positions, ligand_positions))28            if min_dis < 8.0:29                return 130            else:31                return 032    33    io = PDBIO()34    io.set_structure(structure)35    fn = "BS_tmp_"+str(key)+".pdb"36    io.save(fn, ResidueSelect())37    try:38        m2 = Chem.MolFromPDBFile(fn)39        # may contain metal atom, causing MolFromPDBFile return None40        if m2 is None:41            print("first read PDB fail",fn)42            # copy file to tmp dir 43            remove_zn_dir="./docker_result_remove_ZN"44            if not os.path.exists(remove_zn_dir):45                os.mkdir(remove_zn_dir)46            cmd=f"cp {fn}   {remove_zn_dir}"47            print(cmd)48            os.system(cmd)49            fn_remove_zn=os.path.join(remove_zn_dir,fn.replace('.pdb','_remove_ZN.pdb'))50            cmd=f"sed -e '/ZN/d'  {fn}  > {fn_remove_zn}"51            os.system(cmd)52            print("delete metal atom and get new pdb file",fn_remove_zn)53            m2 = Chem.MolFromPDBFile(fn_remove_zn)54        else:55            os.system("rm -f " + fn)56    except:57        print("Read PDB fail for other unknow reason",fn)58    return m259 60def preprocessor(ligand_dir,data_dir):61    """62    get pocket from docking result and save to file:(m1,m2)63 64    input:65        docking_result_sdf_fn: docking result sdf file, one ligand in sdf file will speed up this process in multi-process66        origin_recptor_pdb: receptor pdb file67        data_dir: path for save pocket file68    output:69        0: success70        -1: fail71    """72    file_flag = os.path.basename(ligand_dir)73    print(file_flag)74    try:75        m1 = read_molecule(os.path.join(ligand_dir, f'{file_flag}_ligand.sdf'), remove_hs=True, sanitize=True)76        if m1 is None:  # read mol2 file if sdf file cannot be sanitized77            print('Using the .sdf file failed. We found a .mol2 file instead and are trying to use that.')78            m1 = read_molecule(os.path.join(ligand_dir, f'{file_flag}_ligand.mol2'), remove_hs=True, sanitize=True)79    except Exception as e:80        print(e)81        return -182    if not os.path.exists(data_dir):83        os.mkdir(data_dir)84    if m1 is not None: #docking ligand file may be 0 size85            dst_dir = os.path.join(data_dir,file_flag)86            print(dst_dir)87            if not os.path.exists(dst_dir):88                print('not exist dst dir ,amke it!')89                # dst_dir = os.path.join(data_dir,file_flag)90                os.mkdir(dst_dir)91                ligand_fn = os.path.join(ligand_dir, f'{file_flag}_ligand.sdf')92                print('ligand file ',ligand_fn)93                # dst_dir = os.path.join(data_dir,file_flag)94                os.system(f'cp {ligand_fn} {dst_dir}')95                # os.mkdir(os.path.join(data_dir,file_flag),exist_ok=True)96 97                if len(m1.GetConformers())==0:98                    print(f"{file_flag} mol no conformer!")99                    return -1100                try:101                    pdb_path = os.path.join(ligand_dir, f'{file_flag}_protein.pdb')102                    m2 = extract(m1,pdb_path ,file_flag)103                except:104                    print(f'extract m2 failed {file_flag}')105                    return -1106 107                if m2 is None :108                    print(f"{file_flag} no extracted binding pocket!")109                    # continue110                    return -1111                if len(m2.GetConformers())==0:112                    print(f"{file_flag} receptor no conformer!")113                    return -1114                # save pdb pocket115                Chem.MolToPDBFile(m2, os.path.join(data_dir,file_flag,f'{file_flag}_pocket.pdb'))116 117            else:118                print(f'file done before so skip it  {file_flag}')119                120                return 0121        # return 0122            123    else:124        print("read mol fail")125        return -1126def out_sdf(lig,filename):127    writer = Chem.SDWriter(filename)128    writer.write(lig)129    writer.close()130    return131def get_pocket_with_water(complex_sample):132    status=preprocessor(complex_sample,out_data_dir)133    # print(status)134def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=False):135    if molecule_file.endswith('.mol2'):136        mol = Chem.MolFromMol2File(molecule_file, sanitize=False, removeHs=False)137    elif molecule_file.endswith('.sdf'):138        supplier = Chem.SDMolSupplier(molecule_file, sanitize=False, removeHs=False)139        mol = supplier[0]140    elif molecule_file.endswith('.pdbqt'):141        with open(molecule_file) as file:142            pdbqt_data = file.readlines()143        pdb_block = ''144        for line in pdbqt_data:145            pdb_block += '{}\n'.format(line[:66])146        mol = Chem.MolFromPDBBlock(pdb_block, sanitize=False, removeHs=False)147    elif molecule_file.endswith('.pdb'):148        mol = Chem.MolFromPDBFile(molecule_file, sanitize=False, removeHs=False)149    else:150        raise ValueError('Expect the format of the molecule_file to be '151                         'one of .mol2, .sdf, .pdbqt and .pdb, got {}'.format(molecule_file))152 153    try:154        if sanitize or calc_charges:155            Chem.SanitizeMol(mol)156 157        if calc_charges:158            # Compute Gasteiger charges on the molecule.159            try:160                AllChem.ComputeGasteigerCharges(mol)161            except:162                warnings.warn('Unable to compute charges for the molecule.')163 164        if remove_hs:165            mol = Chem.RemoveHs(mol, sanitize=sanitize)166    except Exception as e:167        print(e)168        print("RDKit was unable to read the molecule.")169        return None170 171    return mol172if __name__ == '__main__':173 174    import time175    from multiprocessing import Pool176    import os177    import gzip178    import tqdm179    # get pocket and save to file180    import argparse181    parser = argparse.ArgumentParser(description='Process data from docking result')182    parser.add_argument("--PDBbind_path", help="file path for save compounds from docking result.", type=str, \183        default='/home/house/caoduanhua/DeepLearningForDock/datasets/dockingModelTestDataset/astex_diverse_set',required=False)184    # parser.add_argument("--docking_result", help="docking result filname.maegz,filename.mae or filename.sdf.", type=str,default=None,required=True)185    # parser.add_argument("--recptor_pdb", help="receptor pdb file.", type=str,default=None,required=True)186    parser.add_argument("--save_dir", help="save pocket file dir.", type=str,default='/home/house/caoduanhua/DeepLearningForDock/datasets/equibind_and_diffdock_dataset/PDBBIND/astex_diverse_set_8A',required=False)187    # parser.add_argument("--prefix", help="Anything that helps you distinguish between compounds.", type=str,default='')188    parser.add_argument("--process_num", help="process num for multi process ", type=int,default=20)189    args = parser.parse_args()190    191    total_sdfs = [os.path.join(args.PDBbind_path,filename) for filename in os.listdir(args.PDBbind_path)]192 193 194    file_tuple_list = []195    for complex_sample in total_sdfs:196        # receptor_fn=args.recptor_pdb197        file_tuple_list.append(complex_sample)198        # print()199    print('num compounds to get pocket',len(file_tuple_list))200    out_data_dir = args.save_dir201    p = Pool(args.process_num)202    pbar = tqdm.tqdm(total=len(file_tuple_list))203    pbar.set_description('get_pocket:')204    update = lambda *args: pbar.update() # set callback function to update pbar state when process end205    for file_tuple in file_tuple_list:206        p.apply_async(get_pocket_with_water,args = (file_tuple,),callback=update)207    print('waiting for processing!')208    p.close()209    p.join()210    print("all pocket done! check the outdir plz!")