| |
| |
| import os |
| import re |
| import argparse |
|
|
| import numpy as np |
|
|
| from utils.logger import print_log |
| from data.converter.blocks_to_data import blocks_to_data |
| from data.converter.mol2_to_blocks import mol2_to_blocks |
| from data.converter.pdb_to_list_blocks import pdb_to_list_blocks |
| from data.mmap_dataset import create_mmap |
|
|
|
|
| def parse(): |
| parser = argparse.ArgumentParser(description='Process PDBBind') |
| parser.add_argument('--data_dir', type=str, required=True, |
| help='Directory of scPDB data') |
| parser.add_argument('--out_dir', type=str, required=True, |
| help='Output directory') |
| return parser.parse_args() |
|
|
|
|
| def parse_index(fpath): |
| with open(fpath, 'r') as fin: |
| lines = fin.readlines() |
| |
| data = {} |
| for line in lines: |
| if line.startswith('#'): |
| continue |
| line = re.split(r'\s+', line) |
| pdb_id, resolution, year, kd = line[:4] |
| data[pdb_id] = kd |
| return data |
|
|
|
|
| def process_iterator_PP(data_dir): |
| indexes = parse_index(os.path.join(data_dir, 'index', 'INDEX_general_PP.2020')) |
| for pdb_id in indexes: |
| list_blocks = pdb_to_list_blocks(os.path.join(data_dir, f'{pdb_id}.ent.pdb')) |
|
|
| data = blocks_to_data(*list_blocks) |
| for key in data: |
| if isinstance(data[key], np.ndarray): |
| data[key] = data[key].tolist() |
|
|
| yield pdb_id, data, [len(data['B']), indexes[pdb_id]] |
|
|
|
|
| def process_iterator_PL(data_dir, index_file): |
| indexes = parse_index(index_file) |
| for pdb_id in indexes: |
| if not os.path.exists(os.path.join(data_dir, pdb_id)): |
| continue |
|
|
| prot_fname = os.path.join(data_dir, pdb_id, f'{pdb_id}_protein.pdb') |
| sm_fname = os.path.join(data_dir, pdb_id, f'{pdb_id}_ligand.mol2') |
|
|
| list_blocks1 = pdb_to_list_blocks(prot_fname) |
| blocks2 = mol2_to_blocks(sm_fname) |
|
|
| data = blocks_to_data(*(list_blocks1 + [blocks2])) |
| for key in data: |
| if isinstance(data[key], np.ndarray): |
| data[key] = data[key].tolist() |
|
|
| yield pdb_id, data, [len(data['B']), indexes[pdb_id]] |
| |
|
|
|
|
|
|
| def main(args): |
| |
| print_log(f'Processing PP') |
| PP_index_file = os.path.join(args.data_dir, 'PP', 'index', 'INDEX_general_PP.2020') |
| PP_index = parse_index(PP_index_file) |
| create_mmap( |
| process_iterator_PP(os.path.join(args.data_dir, 'PP')), |
| os.path.join(args.out_dir, 'PP'), len(PP_index) |
| ) |
|
|
| print_log(f'Processing PL refined set') |
| PL_refine_index_file = os.path.join(args.data_dir, 'refined-set', 'index', 'INDEX_refined_set.2020') |
| PL_refine_index = parse_index(PL_refine_index_file) |
| create_mmap( |
| process_iterator_PL( |
| os.path.join(args.data_dir, 'refined-set'), |
| PL_refine_index_file, |
| ), os.path.join(args.out_dir, 'refined-set'), len(PL_refine_index) |
| ) |
|
|
| print_log(f'Processing PL others') |
| PL_other_index_file = os.path.join(args.data_dir, 'v2020-other-PL', 'index', 'INDEX_general_PL.2020') |
| PL_other_index = parse_index(PL_other_index_file) |
| create_mmap( |
| process_iterator_PL( |
| os.path.join(args.data_dir, 'v2020-other-PL'), |
| PL_other_index_file, |
| ), os.path.join(args.out_dir, 'v2020-other-PL'), len(PL_other_index) - len(PL_refine_index) |
| ) |
|
|
| print_log('Finished!') |
|
|
|
|
| if __name__ == '__main__': |
| main(parse()) |
|
|
