HGNN/align.py

import glob
import skbio
import json, pickle, os
from skbio import alignment
from skbio import Protein
from tqdm import tqdm
from multiprocessing import Pool
import numpy as np
import torch


cutoff = 5.0
blosum50 = \
    {
        '*': {'*': 1, 'A': -5, 'C': -5, 'B': -5, 'E': -5, 'D': -5, 'G': -5,
              'F': -5, 'I': -5, 'H': -5, 'K': -5, 'M': -5, 'L': -5,
              'N': -5, 'Q': -5, 'P': -5, 'S': -5, 'R': -5, 'T': -5,
              'W': -5, 'V': -5, 'Y': -5, 'X': -5, 'Z': -5},
        'A': {'*': -5, 'A': 5, 'C': -1, 'B': -2, 'E': -1, 'D': -2, 'G': 0,
              'F': -3, 'I': -1, 'H': -2, 'K': -1, 'M': -1, 'L': -2,
              'N': -1, 'Q': -1, 'P': -1, 'S': 1, 'R': -2, 'T': 0, 'W': -3,
              'V': 0, 'Y': -2, 'X': -1, 'Z': -1},
        'C': {'*': -5, 'A': -1, 'C': 13, 'B': -3, 'E': -3, 'D': -4,
              'G': -3, 'F': -2, 'I': -2, 'H': -3, 'K': -3, 'M': -2,
              'L': -2, 'N': -2, 'Q': -3, 'P': -4, 'S': -1, 'R': -4,
              'T': -1, 'W': -5, 'V': -1, 'Y': -3, 'X': -1, 'Z': -3},
        'B': {'*': -5, 'A': -2, 'C': -3, 'B': 6, 'E': 1, 'D': 6, 'G': -1,
              'F': -4, 'I': -4, 'H': 0, 'K': 0, 'M': -3, 'L': -4, 'N': 5,
              'Q': 0, 'P': -2, 'S': 0, 'R': -1, 'T': 0, 'W': -5, 'V': -3,
              'Y': -3, 'X': -1, 'Z': 1},
        'E': {'*': -5, 'A': -1, 'C': -3, 'B': 1, 'E': 6, 'D': 2, 'G': -3,
              'F': -3, 'I': -4, 'H': 0, 'K': 1, 'M': -2, 'L': -3, 'N': 0,
              'Q': 2, 'P': -1, 'S': -1, 'R': 0, 'T': -1, 'W': -3, 'V': -3,
              'Y': -2, 'X': -1, 'Z': 5},
        'D': {'*': -5, 'A': -2, 'C': -4, 'B': 6, 'E': 2, 'D': 8, 'G': -1,
              'F': -5, 'I': -4, 'H': -1, 'K': -1, 'M': -4, 'L': -4, 'N': 2,
              'Q': 0, 'P': -1, 'S': 0, 'R': -2, 'T': -1, 'W': -5, 'V': -4,
              'Y': -3, 'X': -1, 'Z': 1},
        'G': {'*': -5, 'A': 0, 'C': -3, 'B': -1, 'E': -3, 'D': -1, 'G': 8,
              'F': -4, 'I': -4, 'H': -2, 'K': -2, 'M': -3, 'L': -4, 'N': 0,
              'Q': -2, 'P': -2, 'S': 0, 'R': -3, 'T': -2, 'W': -3, 'V': -4,
              'Y': -3, 'X': -1, 'Z': -2},
        'F': {'*': -5, 'A': -3, 'C': -2, 'B': -4, 'E': -3, 'D': -5,
              'G': -4, 'F': 8, 'I': 0, 'H': -1, 'K': -4, 'M': 0, 'L': 1,
              'N': -4, 'Q': -4, 'P': -4, 'S': -3, 'R': -3, 'T': -2, 'W': 1,
              'V': -1, 'Y': 4, 'X': -1, 'Z': -4},
        'I': {'*': -5, 'A': -1, 'C': -2, 'B': -4, 'E': -4, 'D': -4,
              'G': -4, 'F': 0, 'I': 5, 'H': -4, 'K': -3, 'M': 2, 'L': 2,
              'N': -3, 'Q': -3, 'P': -3, 'S': -3, 'R': -4, 'T': -1,
              'W': -3, 'V': 4, 'Y': -1, 'X': -1, 'Z': -3},
        'H': {'*': -5, 'A': -2, 'C': -3, 'B': 0, 'E': 0, 'D': -1, 'G': -2,
              'F': -1, 'I': -4, 'H': 10, 'K': 0, 'M': -1, 'L': -3, 'N': 1,
              'Q': 1, 'P': -2, 'S': -1, 'R': 0, 'T': -2, 'W': -3, 'V': -4,
              'Y': 2, 'X': -1, 'Z': 0},
        'K': {'*': -5, 'A': -1, 'C': -3, 'B': 0, 'E': 1, 'D': -1, 'G': -2,
              'F': -4, 'I': -3, 'H': 0, 'K': 6, 'M': -2, 'L': -3, 'N': 0,
              'Q': 2, 'P': -1, 'S': 0, 'R': 3, 'T': -1, 'W': -3, 'V': -3,
              'Y': -2, 'X': -1, 'Z': 1},
        'M': {'*': -5, 'A': -1, 'C': -2, 'B': -3, 'E': -2, 'D': -4,
              'G': -3, 'F': 0, 'I': 2, 'H': -1, 'K': -2, 'M': 7, 'L': 3,
              'N': -2, 'Q': 0, 'P': -3, 'S': -2, 'R': -2, 'T': -1, 'W': -1,
              'V': 1, 'Y': 0, 'X': -1, 'Z': -1},
        'L': {'*': -5, 'A': -2, 'C': -2, 'B': -4, 'E': -3, 'D': -4,
              'G': -4, 'F': 1, 'I': 2, 'H': -3, 'K': -3, 'M': 3, 'L': 5,
              'N': -4, 'Q': -2, 'P': -4, 'S': -3, 'R': -3, 'T': -1,
              'W': -2, 'V': 1, 'Y': -1, 'X': -1, 'Z': -3},
        'N': {'*': -5, 'A': -1, 'C': -2, 'B': 5, 'E': 0, 'D': 2, 'G': 0,
              'F': -4, 'I': -3, 'H': 1, 'K': 0, 'M': -2, 'L': -4, 'N': 7,
              'Q': 0, 'P': -2, 'S': 1, 'R': -1, 'T': 0, 'W': -4, 'V': -3,
              'Y': -2, 'X': -1, 'Z': 0},
        'Q': {'*': -5, 'A': -1, 'C': -3, 'B': 0, 'E': 2, 'D': 0, 'G': -2,
              'F': -4, 'I': -3, 'H': 1, 'K': 2, 'M': 0, 'L': -2, 'N': 0,
              'Q': 7, 'P': -1, 'S': 0, 'R': 1, 'T': -1, 'W': -1, 'V': -3,
              'Y': -1, 'X': -1, 'Z': 4},
        'P': {'*': -5, 'A': -1, 'C': -4, 'B': -2, 'E': -1, 'D': -1,
              'G': -2, 'F': -4, 'I': -3, 'H': -2, 'K': -1, 'M': -3,
              'L': -4, 'N': -2, 'Q': -1, 'P': 10, 'S': -1, 'R': -3,
              'T': -1, 'W': -4, 'V': -3, 'Y': -3, 'X': -1, 'Z': -1},
        'S': {'*': -5, 'A': 1, 'C': -1, 'B': 0, 'E': -1, 'D': 0, 'G': 0,
              'F': -3, 'I': -3, 'H': -1, 'K': 0, 'M': -2, 'L': -3, 'N': 1,
              'Q': 0, 'P': -1, 'S': 5, 'R': -1, 'T': 2, 'W': -4, 'V': -2,
              'Y': -2, 'X': -1, 'Z': 0},
        'R': {'*': -5, 'A': -2, 'C': -4, 'B': -1, 'E': 0, 'D': -2, 'G': -3,
              'F': -3, 'I': -4, 'H': 0, 'K': 3, 'M': -2, 'L': -3, 'N': -1,
              'Q': 1, 'P': -3, 'S': -1, 'R': 7, 'T': -1, 'W': -3, 'V': -3,
              'Y': -1, 'X': -1, 'Z': 0},
        'T': {'*': -5, 'A': 0, 'C': -1, 'B': 0, 'E': -1, 'D': -1, 'G': -2,
              'F': -2, 'I': -1, 'H': -2, 'K': -1, 'M': -1, 'L': -1, 'N': 0,
              'Q': -1, 'P': -1, 'S': 2, 'R': -1, 'T': 5, 'W': -3, 'V': 0,
              'Y': -2, 'X': -1, 'Z': -1},
        'W': {'*': -5, 'A': -3, 'C': -5, 'B': -5, 'E': -3, 'D': -5,
              'G': -3, 'F': 1, 'I': -3, 'H': -3, 'K': -3, 'M': -1, 'L': -2,
              'N': -4, 'Q': -1, 'P': -4, 'S': -4, 'R': -3, 'T': -3,
              'W': 15, 'V': -3, 'Y': 2, 'X': -1, 'Z': -2},
        'V': {'*': -5, 'A': 0, 'C': -1, 'B': -3, 'E': -3, 'D': -4, 'G': -4,
              'F': -1, 'I': 4, 'H': -4, 'K': -3, 'M': 1, 'L': 1, 'N': -3,
              'Q': -3, 'P': -3, 'S': -2, 'R': -3, 'T': 0, 'W': -3, 'V': 5,
              'Y': -1, 'X': -1, 'Z': -3},
        'Y': {'*': -5, 'A': -2, 'C': -3, 'B': -3, 'E': -2, 'D': -3,
              'G': -3, 'F': 4, 'I': -1, 'H': 2, 'K': -2, 'M': 0, 'L': -1,
              'N': -2, 'Q': -1, 'P': -3, 'S': -2, 'R': -1, 'T': -2, 'W': 2,
              'V': -1, 'Y': 8, 'X': -1, 'Z': -2},
        'X': {'*': -5, 'A': -1, 'C': -1, 'B': -1, 'E': -1, 'D': -1,
              'G': -1, 'F': -1, 'I': -1, 'H': -1, 'K': -1, 'M': -1,
              'L': -1, 'N': -1, 'Q': -1, 'P': -1, 'S': -1, 'R': -1,
              'T': -1, 'W': -1, 'V': -1, 'Y': -1, 'X': -1, 'Z': -1},
        'Z': {'*': -5, 'A': -1, 'C': -3, 'B': 1, 'E': 5, 'D': 1, 'G': -2,
              'F': -4, 'I': -3, 'H': 0, 'K': 1, 'M': -1, 'L': -3, 'N': 0,
              'Q': 4, 'P': -1, 'S': 0, 'R': 0, 'T': -1, 'W': -2, 'V': -3,
              'Y': -2, 'X': -1, 'Z': 5}}


import math
def get_align_score(fasta_1, fasta_2):
    kwargs = {}
    kwargs['suppress_sequences'] = False
    kwargs['zero_index'] = True
    kwargs['protein'] = True
    kwargs['substitution_matrix'] = blosum50
    query = alignment.StripedSmithWaterman(fasta_1, **kwargs)
    align = query(fasta_2)
    score = align.optimal_alignment_score
    return float(score)


def read_data(data_root, result_root):
    training_data_fastas = json.load(open(f"{data_root}/align_fastas_dict_10.23.json"))
    bdb_fastas_dict = training_data_fastas['bdb_fastas']
    pdbbind_fastas_dict = training_data_fastas['pdb_fastas']

    save_dir_bdb = f"{result_root}/BDB"
    save_dir_pdbbind = f"{result_root}/PDBBind"
    mol_feat_train_bdb = np.load(f'{save_dir_bdb}/bdb_mol_reps.npy')
    pocket_feat_train_bdb = np.load(f'{save_dir_bdb}/bdb_pocket_reps.npy')
    pocket_names_bdb = json.load(open(f"{save_dir_bdb}/bdb_pocket_names.json"))
    mol_smis_bdb = json.load(open(f"{save_dir_bdb}/bdb_mol_smis.json"))
    bdb_pocket_feat_dict = {pocket_names_bdb[i]: pocket_feat_train_bdb[i] for i in range(len(pocket_names_bdb))}
    bdb_mol_feat_dict = {mol_smis_bdb[i]: mol_feat_train_bdb[i] for i in range(len(mol_smis_bdb))}

    mol_feat_train_pdbbind = np.load(f'{save_dir_pdbbind}/train_mol_reps.npy')
    pocket_feat_train_pdbbind = np.load(f'{save_dir_pdbbind}/train_pocket_reps.npy')
    pocket_names_pdbbind = json.load(open(f"{save_dir_pdbbind}/train_pdbbind_ids.json"))
    mol_smis_pdbbind = json.load(open(f"{save_dir_pdbbind}/train_mol_smis.json"))
    pdbbind_pocket_feat_dict = {pocket_names_pdbbind[i]: pocket_feat_train_pdbbind[i] for i in range(len(pocket_names_pdbbind))}
    pdbbind_mol_feat_dict = {mol_smis_pdbbind[i]: mol_feat_train_pdbbind[i] for i in range(len(mol_smis_pdbbind))}

    return bdb_fastas_dict, bdb_pocket_feat_dict, bdb_mol_feat_dict, pdbbind_fastas_dict, pdbbind_pocket_feat_dict, pdbbind_mol_feat_dict

def get_neighbor_pocket(test_fasta, data_root, result_root, device):
    # 1. Read data file
    print("reading datas")
    bdb_fastas_dict, bdb_pocket_feat_dict, bdb_mol_feat_dict, pdbbind_fastas_dict, \
      pdbbind_pocket_feat_dict, pdbbind_mol_feat_dict = read_data(data_root, result_root)

    training_assay = json.load(open(f"{data_root}/train_label_blend_seq_full.json"))
    training_assay += json.load(open(f"{data_root}/train_label_pdbbind_seq.json"))
    assay_dict = {}
    for assay in training_assay:
        assay["ligands"] = sorted(assay["ligands"], key=lambda x: x["act"], reverse=True)
        if "assay_id" in assay:
            assay_dict[assay["assay_id"]] = assay 
        else:
            assay_dict[assay["pockets"][0][:4]] = assay 

    skip = 0
    # 2. run alignment
    print("running alignment pdbbind")
    align_res_list = []
    for a_name, fasta in tqdm(pdbbind_fastas_dict.items()):
        if a_name not in pdbbind_pocket_feat_dict:
            skip += 1
            continue
        p_name = a_name
        l_smi = assay_dict[a_name]["ligands"][0]["smi"]
        align_score = get_align_score(test_fasta, fasta) / get_align_score(test_fasta, test_fasta)
        if align_score >= 0.5:
            align_res_list.append((pdbbind_pocket_feat_dict[p_name], pdbbind_mol_feat_dict[l_smi], align_score, a_name))

    print("running alignment bindingdb")
    for a_name, fasta in tqdm(bdb_fastas_dict.items()):
        if a_name not in assay_dict:
            skip += 1
            continue
        p_name = assay_dict[a_name]["pockets"][0]
        l_smi = assay_dict[a_name]["ligands"][0]["smi"]
        if l_smi not in bdb_mol_feat_dict:
            continue
        align_score = get_align_score(test_fasta, fasta) / get_align_score(test_fasta, test_fasta)
        if align_score >= 0.5:
            align_res_list.append((bdb_pocket_feat_dict[p_name], bdb_mol_feat_dict[l_smi], align_score, a_name))

    for i, res in enumerate(align_res_list):
        align_res_list[i] = (torch.tensor(res[0]).float().to(device), 
                             torch.tensor(res[1]).float().to(device), 
                             torch.tensor(int((res[2] - 0.5) * 10)).to(device))

    return align_res_list