EPT/scripts/process_data/process_PDBBind_RNA_label.py

#!/usr/bin/python
# -*- coding:utf-8 -*-
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
from data.format import VOCAB, Atom, Block

import copy


import math

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 kd_to_dg(kd, temperature=25.0):
    """Conversion of Kd to DG"""
    R = 0.0019872043
    dg_rt = math.log(kd)
    temp_in_k = temperature + 273.15
    rt = R * temp_in_k
    return dg_rt * rt

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

        if (not kd.startswith('Kd')) and (not kd.startswith('Ki')):  # IC50 is very different from Kd and Ki, therefore discarded
            print_log(f'{pdb_id} not measured by Kd or Ki, dropped.', level='ERROR')
            # return None
            continue
        
        if '=' not in kd:  # some data only provide a threshold, e.g. Kd<1nM, discarded
            print_log(f'{pdb_id} Kd only has threshold: {kd}', level='ERROR')
            # return None
            continue

        kd = kd.split('=')[-1].strip()
        aff, unit = float(kd[:-2]), kd[-2:]
        if unit == 'mM':
            aff *= 1e-3
        elif unit == 'nM':
            aff *= 1e-9
        elif unit == 'uM':
            aff *= 1e-6
        elif unit == 'pM':
            aff *= 1e-12
        elif unit == 'fM':
            aff *= 1e-15
        else:
            # return None   # unrecognizable unit
            continue
        
        # affinity data
        data[pdb_id] = {
            'Kd': aff,
            'dG': kd_to_dg(aff, 25.0),   # regard as measured under the standard condition
            'neglog_aff': -math.log(aff, 10)  # pK = -log_10 (Kd)
        }


    return data



def blocks_to_coords(blocks: List[Block]):
    max_n_unit = 0
    coords, masks = [], []
    for block in blocks:
        coords.append([unit.get_coord() for unit in block.units])
        max_n_unit = max(max_n_unit, len(coords[-1]))
        masks.append([1 for _ in coords[-1]])
    
    for i in range(len(coords)):
        num_pad =  max_n_unit - len(coords[i])
        coords[i] = coords[i] + [[0, 0, 0] for _ in range(num_pad)]
        masks[i] = masks[i] + [0 for _ in range(num_pad)]
    
    return np.array(coords), np.array(masks).astype('bool')  # [N, M, 3], [N, M], M == max_n_unit, in mask 0 is for padding


def dist_matrix_from_coords(coords1, masks1, coords2, masks2):
    dist = np.linalg.norm(coords1[:, None] - coords2[None, :], axis=-1)  # [N1, N2, M]
    dist = dist + np.logical_not(masks1[:, None] * masks2[None, :]) * 1e6  # [N1, N2, M]
    dist = np.min(dist, axis=-1)  # [N1, N2]
    return dist


def dist_matrix_from_residues(residue_list1, residue_list2):
    coords, mask = blocks_to_coords(residue_list1 + residue_list2)
    midpoint = len(residue_list1)
    coords1, masks1 = coords[:midpoint], mask[:midpoint]
    coords2, masks2 = coords[midpoint:], mask[midpoint:]
    return dist_matrix_from_coords(coords1, masks1, coords2, masks2)

def blocks_interface(blocks1, blocks2, dist_th):
    blocks_coord, blocks_mask = blocks_to_coords(blocks1 + blocks2)
    blocks1_coord, blocks1_mask = blocks_coord[:len(blocks1)], blocks_mask[:len(blocks1)]
    blocks2_coord, blocks2_mask = blocks_coord[len(blocks1):], blocks_mask[len(blocks1):]
    dist = dist_matrix_from_coords(blocks1_coord, blocks1_mask, blocks2_coord, blocks2_mask)
    
    on_interface = dist < dist_th
    indexes1 = np.nonzero(on_interface.sum(axis=1) > 0)[0]
    indexes2 = np.nonzero(on_interface.sum(axis=0) > 0)[0]

    blocks1 = [blocks1[i] for i in indexes1]
    blocks2 = [blocks2[i] for i in indexes2]

    return blocks1, blocks2

def break_blocks_into_atoms(blocks):

    block_list = []

    for block in blocks:
        for atom in block.units:
            atom_new = copy.deepcopy(atom)
            atom_new.pos_code = VOCAB.atom_pos_sm
            block_list.append(
                Block(symbol=atom_new.element.lower(), units = [atom_new])
            )

    return block_list


def process_iterator_PP(data_dir, index_file, dist_th = 6.0):
    indexes = parse_index(index_file)
    for pdb_id in indexes:
        list_blocks, chains = pdb_to_list_blocks(os.path.join(data_dir, f'{pdb_id}.ent.pdb'), return_chain_ids=True)

        if len(list_blocks) != 2:
            continue

        rec_residues, lig_residues = list_blocks # [Q] I can't visit FASTA. Swapping the order of the two chains may not actually change the logic of the following codes?

        rec_inter, lig_inter = blocks_interface(rec_residues, lig_residues, dist_th)

        list_blocks = [rec_inter, lig_inter]


        data = blocks_to_data(*list_blocks)
        for key in data:
            if isinstance(data[key], np.ndarray):
                data[key] = data[key].tolist()

        Kd, dG, neglog_aff = indexes[pdb_id]['Kd'], indexes[pdb_id]['dG'], indexes[pdb_id]['neglog_aff']

        yield pdb_id, data, [len(data['B']), Kd, dG, neglog_aff]


def process_iterator_PL(data_dir, index_file, dist_th = 6.0):
    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)

        blocks1 = []
        for b in list_blocks1:
            blocks1.extend(b)

        # construct pockets
        blocks1, _ = blocks_interface(blocks1, blocks2, dist_th)

        list_blocks = [blocks1, blocks2]

        data = blocks_to_data(*list_blocks)
        for key in data:
            if isinstance(data[key], np.ndarray):
                data[key] = data[key].tolist()


        Kd, dG, neglog_aff = indexes[pdb_id]['Kd'], indexes[pdb_id]['dG'], indexes[pdb_id]['neglog_aff']

        yield pdb_id, data, [len(data['B']), Kd, dG, neglog_aff]

def process_iterator_NL(data_dir, index_file, dist_th = 6.0):
    indexes = parse_index(index_file)
    for pdb_id in indexes:
        if not os.path.exists(os.path.join(data_dir, pdb_id)):
            continue

        list_blocks, chains = pdb_to_list_blocks(os.path.join(data_dir, f'{pdb_id}.ent.pdb'), return_chain_ids=True)

        bases = ['DA', 'DG', 'DC', 'DT', 'R-A', 'R-G', 'R-C', 'R-U']

        rec_blocks = []
        lig_blocks = []
        for chain in list_blocks:
            split_point = None
            for i in range(len(chain)):
                residue = chain[len(chain) - i - 1]
                if residue.symbol in bases:
                    split_point = len(chain) - i
                    break
            rec_blocks.extend(chain[:split_point])
            lig_blocks.extend(break_blocks_into_atoms(chain[split_point:]))

        # construct pockets
        rec_interface, _ = blocks_interface(rec_blocks, lig_blocks, dist_th)

        list_blocks = [rec_interface, lig_blocks]

        data = blocks_to_data(*list_blocks)
        for key in data:
            if isinstance(data[key], np.ndarray):
                data[key] = data[key].tolist()


        Kd, dG, neglog_aff = indexes[pdb_id]['Kd'], indexes[pdb_id]['dG'], indexes[pdb_id]['neglog_aff']

        yield pdb_id, data, [len(data['B']), Kd, dG, neglog_aff]

        



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')),
        PP_index_file,
        os.path.join(args.out_dir, 'PP-aff'), 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-aff'), 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-aff'), len(PL_other_index) - len(PL_refine_index)
    )

    print_log(f'Processing PN')
    PN_index_file = os.path.join(args.data_dir, 'PN', 'index', 'INDEX_general_PN.2020')
    PN_index = parse_index(PN_index_file)
    create_mmap(
        process_iterator_PP(os.path.join(args.data_dir, 'PN')),
        PN_index_file,
        os.path.join(args.out_dir, 'PN-aff'), len(PN_index)
    )

    print_log(f'Processing NL')
    NL_index_file = os.path.join(args.data_dir, 'NL', 'index', 'INDEX_general_NL.2020')
    NL_index = parse_index(PP_index_file)
    create_mmap(
        process_iterator_NL(os.path.join(args.data_dir, 'NL')),
        NL_index_file,
        os.path.join(args.out_dir, 'NL-aff'), len(NL_index)
    )

    print_log('Finished!')


if __name__ == '__main__':
    main(parse())