File size: 4,643 Bytes
4947683 | 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 | #!/usr/bin/python
# -*- coding:utf-8 -*-
import os
import sys
import json
import pickle
import argparse
import numpy as np
PROJ_DIR = os.path.join(
os.path.split(os.path.abspath(__file__))[0],
'..', '..'
)
print(f'Project directory: {PROJ_DIR}')
sys.path.append(PROJ_DIR)
from utils.logger import print_log
from data.converter.pdb_to_list_blocks import pdb_to_list_blocks
from data.converter.mol2_to_blocks import mol2_to_blocks
from data.converter.blocks_interface import blocks_interface
from data.converter.blocks_to_data import blocks_to_data
from data.mmap_dataset import create_mmap
def parse():
parser = argparse.ArgumentParser(description='Process PDBbind benchmark of protein-ligand interaction')
parser.add_argument('--benchmark_dir', type=str, required=True,
help='Directory of the benchmark containing metadata and pdb_files')
parser.add_argument('--out_dir', type=str, required=True,
help='Output directory')
parser.add_argument('--interface_dist_th', type=float, default=8.0,
help='Residues who has atoms with distance below this threshold are considered in the complex interface')
return parser.parse_args()
def process_iterator(benchmark_dir, interface_dist_th):
labels = json.load(open(os.path.join(benchmark_dir, 'metadata', 'affinities.json'), 'r'))
for pdb_id in labels:
pdb_dir = os.path.join(benchmark_dir, 'pdb_files')
prot_fname = os.path.join(pdb_dir, pdb_id, pdb_id + '.pdb')
sm_fname = os.path.join(pdb_dir, pdb_id, f'{pdb_id}_ligand.mol2')
list_blocks1 = pdb_to_list_blocks(prot_fname)
blocks2 = mol2_to_blocks(sm_fname)
# try:
# list_blocks1 = pdb_to_list_blocks(prot_fname)
# except Exception as e:
# print_log(f'{pdb_id} protein parsing failed: {e}', level='ERROR')
# continue
# try:
# blocks2 = mol2_to_blocks(sm_fname)
# except Exception as e:
# print_log(f'{pdb_id} ligand parsing failed: {e}', level='ERROR')
# continue
blocks1 = []
for b in list_blocks1:
blocks1.extend(b)
data = blocks_to_data(blocks1, blocks2)
for key in data:
if isinstance(data[key], np.ndarray):
data[key] = data[key].tolist()
# construct pockets
blocks1, _ = blocks_interface(blocks1, blocks2, interface_dist_th)
if len(blocks1) == 0: # no interface (if len(interface1) == 0 then we must have len(interface2) == 0)
print_log(f'{pdb_id} has no interface', level='ERROR')
continue
data_interface = blocks_to_data(blocks1, blocks2)
for key in data_interface:
if isinstance(data_interface[key], np.ndarray):
data_interface[key] = data_interface[key].tolist()
result = {
'complex': data,
'interface': data_interface
}
yield pdb_id, result, [len(data['B']), len(data_interface['B']), labels[pdb_id]]
def main(args):
# TODO: 1. preprocess PDBbind into json summaries and complex pdbs
labels = json.load(open(os.path.join(args.benchmark_dir, 'metadata', 'affinities.json'), 'r'))
print_log(f'Processing data from directory: {args.benchmark_dir}.')
create_mmap(
process_iterator(args.benchmark_dir, args.interface_dist_th),
args.out_dir, len(labels))
print_log('Finished database construction!')
id2line = {}
with open(os.path.join(args.out_dir, 'index.txt'), 'r') as fin:
for line in fin.readlines():
_id = line.split('\t')[0]
id2line[_id] = line
for split in ['identity30', 'identity60', 'scaffold']:
split_info = json.load(open(os.path.join(args.benchmark_dir, 'metadata', f'{split}_split.json'), 'r'))
out_dir = os.path.join(args.out_dir, split)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
for name in ['train', 'valid', 'test']:
data_out_path = os.path.join(out_dir, name + '.txt')
data_out = []
miss_cnt = 0
for pdb_id in split_info[name]:
if pdb_id in id2line:
data_out.append(id2line[pdb_id])
else:
miss_cnt += 1
print_log(f'Obtained {len(data_out)}, missing {miss_cnt}, saving to {data_out_path}...')
with open(data_out_path, 'w') as fout:
fout.writelines(data_out)
print_log('Finished splitting!')
if __name__ == '__main__':
main(parse())
|