File size: 3,826 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 | import pandas as pd
# from defaultdict import defaultdict
from collections import defaultdict
import os
from argparse import ArgumentParser, Namespace, FileType
parser = ArgumentParser()
parser.add_argument('--data_dir', type=str, default='~/SurfDock/model/data/test_samples', help='')
parser.add_argument('--surface_out_dir', type=str, default='~/SurfDock/model/data/test_samples_8A_surface', help='')
parser.add_argument('--Screen_ligand_library_file', type=str, default=None, help='')
parser.add_argument('--output_csv_file', type=str, default='~/SurfDock/model/data/test_samples_8A_surface', help='')
parser.add_argument('--is_docking_result_dir', action='store_true', default=False, help='')
parser.add_argument('--docking_result_dir', type=str, default='', help='')
# dirname = os.path.splitext(pocket_path.split('/')[-1])[0] + '_'+ os.path.splitext(ligands_path.split('/')[-1])[0]
# write_dir = os.path.join(args.out_dir,'SurfDock_docking_result',dirname)#f'{args.out_dir}/SurfDock_docking_result/{dirname}'
args = parser.parse_args()
os.makedirs(os.path.dirname(args.output_csv_file),exist_ok=True)
from tqdm import tqdm
args_list=defaultdict(list)
proteins = [i for i in os.listdir(args.surface_out_dir) if os.path.isdir(os.path.join(args.surface_out_dir, i)) ]
for protein in tqdm(proteins ):
target_filename = os.path.join(args.surface_out_dir,protein,f'{protein}_protein_processed_obabel_reduce_obabel.pdb')
if not os.path.exists(target_filename):
target_filename = os.path.join(args.data_dir,protein,f'{protein}_protein_processed.pdb')
if not os.path.exists(target_filename):
raise ValueError(f'{target_filename} not exists , Please check file name or path')
ref_ligand_filename = os.path.join(args.data_dir,protein,f'{protein}_ligand.sdf')
ligand_filename = os.path.join(args.data_dir,protein,f'{protein}_ligand.sdf')
if args.Screen_ligand_library_file is not None:
print(f'Using Screen ligands library file: {args.Screen_ligand_library_file}')
ligand_filename = args.Screen_ligand_library_file
if os.path.exists(ref_ligand_filename):
pocket = os.path.join(args.surface_out_dir, protein, f'{protein}_protein_processed_obabel_reduce_obabel_8A.pdb')
surface = os.path.join(args.surface_out_dir, protein, f'{protein}_protein_processed_obabel_reduce_obabel_8A.ply')
if not os.path.exists(pocket):
pocket = os.path.join(args.surface_out_dir, protein, f'{protein}_protein_processed_8A.pdb')
if not os.path.exists(surface):
surface = os.path.join(args.surface_out_dir, protein, f'{protein}_protein_processed_8A.ply')
if os.path.exists(pocket) and os.path.exists(surface):
args_list['protein_path'].append(target_filename)
args_list['pocket_path'].append(pocket)
args_list['ref_ligand'].append(ref_ligand_filename)
if args.is_docking_result_dir:
dirname = os.path.splitext(pocket.split('/')[-1])[0] + '_'+ os.path.splitext(ligand_filename.split('/')[-1])[0]
# write_dir = os.path.join(args.docking_result_dir,'SurfDock_docking_result',dirname)#f'{args.out_dir}/SurfDock_docking_result/{dirname}'
args_list['ligand_path'].append(os.path.join(args.docking_result_dir,'SurfDock_docking_result',dirname))
else:
args_list['ligand_path'].append(ligand_filename)
args_list['protein_surface'].append(surface)
else:
pass
print(pocket)
else:
print(protein)
pd.DataFrame(args_list).to_csv(args.output_csv_file,index=False)
|