File size: 8,653 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 | #!/usr/bin/python
# -*- coding:utf-8 -*-
import os
import re
import json
import shutil
import argparse
from functools import partial
import numpy as np
from rdkit import Chem
from rdkit.Chem.Scaffolds import MurckoScaffold
from utils.logger import print_log
from data.converter.blocks_to_data import blocks_to_data
from data.converter.sdf_to_list_blocks import sdf_to_list_blocks
from data.converter.pdb_to_list_blocks import pdb_to_list_blocks
from data.converter.blocks_interface import blocks_interface
from data.mmap_dataset import create_mmap
from .non_redundant_pdb import clustering
def parse():
parser = argparse.ArgumentParser(description='Process PDBBind')
parser.add_argument('--data_dir', type=str, required=True,
help='Directory of raw data of general set and refined set')
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 parse_index(fpath, quality):
pdb_dir = os.path.dirname(fpath)
with open(fpath, 'r') as fin:
lines = fin.readlines()
data = {}
for line in lines:
name, status = line.strip().split('\t')
if status == 'failed':
continue
metadata = json.load(open(os.path.join(pdb_dir, name, 'metadata.json'), 'r'))
if 'pos' not in metadata['dock']:
continue # the positive sample failed to dock
metadata['quality'] = quality
data[name] = (metadata, os.path.join(pdb_dir, name))
return data
def process_iterator_PL(indexes, if_th):
for pdb_id in indexes:
metadata, data_dir = indexes[pdb_id]
prot_fname = os.path.join(data_dir, 'receptor.pdb')
sm_fname = os.path.join(data_dir, 'ligands.sdf')
prot_list_blocks = pdb_to_list_blocks(prot_fname)
sm_dicts = sdf_to_list_blocks(sm_fname, dict_form=True, silent=True)
rec_blocks = []
for blocks in prot_list_blocks:
rec_blocks.extend(blocks)
all_data, len_dict = {}, {}
for name in sorted(list(sm_dicts.keys())):
pocket_blocks, _ = blocks_interface(rec_blocks, sm_dicts[name], if_th)
if len(pocket_blocks) == 0:
print_log(f'{pdb_id}, {name} no interaction detected', level='WARN')
data = blocks_to_data(pocket_blocks, sm_dicts[name])
for key in data:
if isinstance(data[key], np.ndarray):
data[key] = data[key].tolist()
all_data[name] = data
len_dict[name] = len(pocket_blocks) + len(sm_dicts[name])
yield pdb_id, all_data, [len_dict, metadata]
def _extract_id_to_seqs(lines, filter_refined=False):
id2seqs = {}
for line in lines:
line = line.strip().split('\t')
prop = json.loads(line[-1])
if prop['quality'] == 'general' and filter_refined:
continue
rec_seq = 'X'.join(prop['receptor_seqs'])
id2seqs[line[0]] = (rec_seq, prop['dock']['pos']['smiles'])
return id2seqs
def _cluster_seq_id(id2seqs, seq_id):
# make temporary directory
tmp_dir = './tmp'
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
else:
raise ValueError(f'Working directory {tmp_dir} exists!')
# 1. get non-redundant dimer by 90% seq-id
fasta = os.path.join(tmp_dir, 'seq.fasta')
# receptor
with open(fasta, 'w') as fout:
for _id in id2seqs:
fout.write(f'>{_id}\n{id2seqs[_id][0]}\n')
id2clu, clu2id = clustering(fasta, tmp_dir, seq_id)
shutil.rmtree(tmp_dir)
return id2clu, clu2id
def _generate_scaffold(smiles, include_chirality=False):
"""return scaffold string of target molecule"""
mol = Chem.MolFromSmiles(smiles)
scaffold = MurckoScaffold\
.MurckoScaffoldSmiles(mol=mol, includeChirality=include_chirality)
return scaffold
def _cluster_scaffold(id2seqs):
id2clu, clu2id = {}, {}
for _id in id2seqs:
smi = id2seqs[_id][1]
scaffold = _generate_scaffold(smi, include_chirality=True)
id2clu[_id] = scaffold
if scaffold not in clu2id:
clu2id[scaffold] = []
clu2id[scaffold].append(_id)
return id2clu, clu2id
def create_test_set(mmap_dir, max_size=500):
index_file = os.path.join(mmap_dir, 'index.txt')
with open(index_file, 'r') as fin:
lines = fin.readlines()
id2lines = { line.split('\t')[0]: line for line in lines }
id2seqs = _extract_id_to_seqs(lines, filter_refined=True) # only use high-quality data for testing
max_cluster_size = 5 # use
# delete similar recptors (above 30% sequence identity)
id2clu_seq, clu2id_seq = _cluster_seq_id(id2seqs, 0.3)
# delete similar scaffolds
id2clu_scaffold, clu2id_scaffold = _cluster_scaffold(id2seqs)
test_ids = []
for _id in id2seqs:
clu_size1 = len(clu2id_seq[id2clu_seq[_id]])
clu_size2 = len(clu2id_scaffold[id2clu_scaffold[_id]])
if clu_size1 < max_cluster_size and clu_size2 < max_cluster_size:
test_ids.append(_id)
np.random.shuffle(test_ids)
test_ids = test_ids[:max_size]
# write results
with open(os.path.join(mmap_dir, 'test.txt'), 'w') as fout:
for _id in test_ids:
fout.write(id2lines[_id])
return test_ids
def _split_by_cluster(id2clu, clu2id, test_ids, train_ratio):
test_clus = { id2clu[_id]: True for _id in test_ids }
available_clus = sorted([ c for c in clu2id if c not in test_clus ])
train_size = int(len(available_clus) * train_ratio)
train_ids, valid_ids = [], []
np.random.shuffle(available_clus)
for c in available_clus[:train_size]:
train_ids.extend(clu2id[c])
for c in available_clus[train_size:]:
valid_ids.extend(clu2id[c])
return train_ids, valid_ids, available_clus[:train_size], available_clus[train_size:]
def split_by_func(mmap_dir, test_ids, name, func, train_ratio=0.9):
index_file = os.path.join(mmap_dir, 'index.txt')
with open(index_file, 'r') as fin:
lines = fin.readlines()
id2lines = { line.split('\t')[0]: line for line in lines }
id2seqs = _extract_id_to_seqs(lines, filter_refined=False)
# cluster
id2clu, clu2id = func(id2seqs)
# split
train_ids, valid_ids, train_clus, valid_clus = _split_by_cluster(id2clu, clu2id, test_ids, train_ratio)
print_log(f'Train set: {len(train_ids)} entries, {len(train_clus)} clusters')
print_log(f'Validation set: {len(valid_ids)} entries, {len(valid_clus)} clusters')
# write results
out_dir = os.path.join(mmap_dir, name)
os.makedirs(out_dir, exist_ok=True)
with open(os.path.join(out_dir, 'train.txt'), 'w') as fout:
for _id in train_ids: fout.write(id2lines[_id])
with open(os.path.join(out_dir, 'train_cluster.txt'), 'w') as fout:
for _id in train_ids: fout.write(f'{_id}\t{id2clu[_id]}\n')
with open(os.path.join(out_dir, 'valid.txt'), 'w') as fout:
for _id in valid_ids: fout.write(id2lines[_id])
with open(os.path.join(out_dir, 'valid_cluster.txt'), 'w') as fout:
for _id in valid_ids: fout.write(f'{_id}\t{id2clu[_id]}\n')
def main(args):
if not os.path.exists(args.out_dir):
print_log(f'Generating data from {args.data_dir}')
# refined set
indexes = parse_index(os.path.join(args.data_dir, 'processed_refined_set', 'done.log'), 'refined')
indexes2 = parse_index(os.path.join(args.data_dir, 'processed_general_set', 'done.log'), 'general')
for name in indexes2:
assert name not in indexes, name
indexes[name] = indexes2[name]
create_mmap(
process_iterator_PL(
indexes, args.interface_dist_th
), args.out_dir, len(indexes)
)
# create splits
non_redundant_test = create_test_set(args.out_dir)
print_log(f'Size of non-redundant test set: {len(non_redundant_test)}')
split_funcs = {
'seqid_30': partial(_cluster_seq_id, seq_id=0.3),
'seqid_60': partial(_cluster_seq_id, seq_id=0.6),
'scaffold': _cluster_scaffold,
}
for name in split_funcs:
print()
print_log(f'Processing split {name}...')
split_by_func(args.out_dir, non_redundant_test, name, split_funcs[name])
print_log('Finished!')
if __name__ == '__main__':
np.random.seed(12)
main(parse()) |