File size: 3,376 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 | #!/usr/bin/python
# -*- coding:utf-8 -*-
import os
import argparse
import numpy as np
import pandas as pd
from rdkit import Chem
import json
import pickle
from utils.logger import print_log
from data.format import Block, Atom, VOCAB
from data.converter.df_to_blocks import df_to_blocks
from data.converter.blocks_to_data import blocks_to_data
from data.converter.blocks_interface import blocks_interface
from data.mmap_dataset import create_mmap
from data.atom3d_lmdb import LMDBDataset
import pdb
def parse():
parser = argparse.ArgumentParser(
description="Process molecule data from MSP dataset."
)
parser.add_argument(
"--base_path", type=str, required=True, help="Directory of dataset"
)
parser.add_argument("--out_dir", type=str, required=True, help="Output directory")
return parser.parse_args()
def process_iterator(base_data):
for item in base_data:
chains = item['original_atoms'].chain.unique()
blocks_original, chain_ids_original = df_to_blocks(item["original_atoms"], key_atom_name="name", return_chain_ids=True)
assert len(blocks_original) == len(chain_ids_original)
blocks_mutated, chain_ids_mutated = df_to_blocks(item["mutated_atoms"], key_atom_name="name", return_chain_ids=True)
assert len(blocks_mutated) == len(chain_ids_mutated)
assert len(blocks_original) == len(blocks_mutated)
ori_seq = ''.join([block.symbol for block in blocks_original])
mut_seq = ''.join([block.symbol for block in blocks_mutated])
is_mut = [0 if a == b else 1 for a, b in zip(ori_seq, mut_seq)]
assert sum(is_mut) == 1
mut_pos = is_mut.index(1)
_, select_index = blocks_interface([blocks_original[mut_pos]], blocks_original, 8.0, return_index=True)
chain2blocks_original = { c: [] for c in chains }
for i in select_index:
chain2blocks_original[chain_ids_original[i]].append(blocks_original[i])
chain2blocks_mutated = { c: [] for c in chains }
for i in select_index:
chain2blocks_mutated[chain_ids_mutated[i]].append(blocks_mutated[i])
label = item["label"] # '1' for better binding, '0' for worse or equal
list_blocks_original = [chain2blocks_original[c] for c in chain2blocks_original if len(chain2blocks_original[c]) > 0]
data1 = blocks_to_data(*list_blocks_original)
list_blocks_mutated = [chain2blocks_mutated[c] for c in chain2blocks_mutated if len(chain2blocks_mutated[c]) > 0]
data2 = blocks_to_data(*list_blocks_mutated)
for key in data1:
if isinstance(data1[key], np.ndarray):
data1[key] = data1[key].tolist()
for key in data2:
if isinstance(data2[key], np.ndarray):
data2[key] = data2[key].tolist()
data = (data1, data2)
# id, data, label
yield item["id"], data, [len(data1['B']) + len(data2['B']), label]
def main(args):
for split in ["train", "val", "test"]:
print_log(f"Processing {split} dataset ...")
base_data = LMDBDataset(os.path.join(args.base_path, split))
create_mmap(
process_iterator(base_data),
os.path.join(args.out_dir, split),
len(base_data),
)
print_log("Finished!")
if __name__ == "__main__":
main(parse())
|