OneScience-Group/SurfDock

1import pandas as pd2# from defaultdict import defaultdict3from collections import defaultdict4import os5from argparse import ArgumentParser, Namespace, FileType6parser = ArgumentParser()7parser.add_argument('--data_dir', type=str, default='~/SurfDock/model/data/test_samples', help='')8parser.add_argument('--surface_out_dir', type=str, default='~/SurfDock/model/data/test_samples_8A_surface', help='')9parser.add_argument('--Screen_ligand_library_file', type=str, default=None, help='')10parser.add_argument('--output_csv_file', type=str, default='~/SurfDock/model/data/test_samples_8A_surface', help='')11parser.add_argument('--is_docking_result_dir', action='store_true', default=False, help='')12parser.add_argument('--docking_result_dir', type=str, default='', help='')13# dirname = os.path.splitext(pocket_path.split('/')[-1])[0] + '_'+ os.path.splitext(ligands_path.split('/')[-1])[0] 14# write_dir =  os.path.join(args.out_dir,'SurfDock_docking_result',dirname)#f'{args.out_dir}/SurfDock_docking_result/{dirname}'15args = parser.parse_args()16 17os.makedirs(os.path.dirname(args.output_csv_file),exist_ok=True)18from tqdm import tqdm19 20args_list=defaultdict(list)21proteins = [i for i in os.listdir(args.surface_out_dir) if os.path.isdir(os.path.join(args.surface_out_dir, i)) ]22for protein in tqdm(proteins ):23        target_filename = os.path.join(args.surface_out_dir,protein,f'{protein}_protein_processed_obabel_reduce_obabel.pdb')24        if not os.path.exists(target_filename):25            target_filename = os.path.join(args.data_dir,protein,f'{protein}_protein_processed.pdb')26        if not os.path.exists(target_filename):27            raise ValueError(f'{target_filename} not exists , Please check file name or path')28 29        ref_ligand_filename = os.path.join(args.data_dir,protein,f'{protein}_ligand.sdf')30        ligand_filename = os.path.join(args.data_dir,protein,f'{protein}_ligand.sdf')31        if args.Screen_ligand_library_file is not None:32            print(f'Using Screen ligands library file: {args.Screen_ligand_library_file}')33            ligand_filename = args.Screen_ligand_library_file34        35        if os.path.exists(ref_ligand_filename):36        37            pocket = os.path.join(args.surface_out_dir, protein, f'{protein}_protein_processed_obabel_reduce_obabel_8A.pdb')38            surface = os.path.join(args.surface_out_dir, protein, f'{protein}_protein_processed_obabel_reduce_obabel_8A.ply')39            if not os.path.exists(pocket):40                pocket = os.path.join(args.surface_out_dir, protein, f'{protein}_protein_processed_8A.pdb')41            if not os.path.exists(surface):42                surface = os.path.join(args.surface_out_dir, protein, f'{protein}_protein_processed_8A.ply')43                44            if os.path.exists(pocket) and os.path.exists(surface):45                args_list['protein_path'].append(target_filename)46                args_list['pocket_path'].append(pocket)47                args_list['ref_ligand'].append(ref_ligand_filename)48                49                if args.is_docking_result_dir:50                    dirname = os.path.splitext(pocket.split('/')[-1])[0] + '_'+ os.path.splitext(ligand_filename.split('/')[-1])[0] 51                    # write_dir =  os.path.join(args.docking_result_dir,'SurfDock_docking_result',dirname)#f'{args.out_dir}/SurfDock_docking_result/{dirname}'52                    args_list['ligand_path'].append(os.path.join(args.docking_result_dir,'SurfDock_docking_result',dirname))53                else:54                    args_list['ligand_path'].append(ligand_filename)55                args_list['protein_surface'].append(surface)56            else:57                pass58                print(pocket)59        else:60            61            print(protein)62pd.DataFrame(args_list).to_csv(args.output_csv_file,index=False)63