HHBLITS氨基酸编号映射到AATYPE编号

HHblits（HMM-HMM比对）是一种用于蛋白质序列比对的工具，它使用了profile-profile比对的方法。在HHblits的输出结果中，通常会包含氨基酸的编号。 蛋白质序列特征 features['aatype']中氨基酸的编号为restypes_with_x_and_gap（热独转化）而多序列比对特征features['msa']特征以及模版特征features['template_aatype']中氨基酸编号用的HHBLITS_AA_TO_ID，因此需要进行转化。

import tensorflow as tf
import pickle
import numpy as np

# This is the standard residue order when coding AA type as a number.
# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
restypes = [
    'A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P',
    'S', 'T', 'W', 'Y', 'V'
]

restypes_with_x_and_gap = restypes + ['X', '-']


# Partial inversion of HHBLITS_AA_TO_ID.
ID_TO_HHBLITS_AA = {
    0: 'A',
    1: 'C',  # Also U.
    2: 'D',  # Also B.
    3: 'E',  # Also Z.
    4: 'F',
    5: 'G',
    6: 'H',
    7: 'I',
    8: 'K',
    9: 'L',
    10: 'M',
    11: 'N',
    12: 'P',
    13: 'Q',
    14: 'R',
    15: 'S',
    16: 'T',
    17: 'V',
    18: 'W',
    19: 'Y',
    20: 'X',  # Includes J and O.
    21: '-',
}


## 映射函数
def correct_msa_restypes(protein):
    """Correct MSA restype to have the same order as residue_constants."""
    new_order_list = MAP_HHBLITS_AATYPE_TO_OUR_AATYPE
    new_order = tf.constant(new_order_list, dtype=protein['msa'].dtype)
  
    # 就是抽取出params的第axis维度上在indices里面所有的index
    protein['msa'] = tf.gather(params=new_order, indices=protein['msa'], axis=0)
    # protein['msa'] = tf.gather(new_order, protein['msa'], axis=0)

    # HHBLITS_AATYPE热独编码转restype的热独编码的矩阵
    perm_matrix = np.zeros((22, 22), dtype=np.float32) 
    perm_matrix[range(len(new_order_list)), new_order_list] = 1.

    for k in protein:
        if 'profile' in k:  # Include both hhblits and psiblast profiles
            num_dim = protein[k].shape.as_list()[-1]
            assert num_dim in [20, 21, 22], (
                'num_dim for %s out of expected range: %s' % (k, num_dim))
            protein[k] = tf.tensordot(protein[k], perm_matrix[:num_dim, :num_dim], 1)
    
    return protein


def fix_templates_aatype(protein):
    """Fixes aatype encoding of templates."""
    # Map one-hot to indices.
    protein['template_aatype'] = tf.argmax(
        protein['template_aatype'], output_type=tf.int32, axis=-1)
    # Map hhsearch-aatype to our aatype.
    new_order_list = MAP_HHBLITS_AATYPE_TO_OUR_AATYPE
    new_order = tf.constant(new_order_list, dtype=tf.int32)
    protein['template_aatype'] = tf.gather(params=new_order,
                                           indices=protein['template_aatype'])
    # 可以再转化会one-hot编码
    return protein


# ID_TO_HHBLITS_AA中的id（key, 0-21）对应的氨基酸在restypes中的索引（index）
# 注：msa特征中features['msa']中氨基酸的编号用的是HHBLITS_AA_TO_ID，
# 而序列特征features['aatype']用的是restypes_with_x_and_gap编号。

# Map hhsearch-aatype to our aatype.
MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
    restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
    for i in range(len(restypes_with_x_and_gap)))

print(MAP_HHBLITS_AATYPE_TO_OUR_AATYPE)

# 读取原始的 HBB_features_lst.pkl
with open("HBB_features_lst.pkl",'rb') as f:
    HBB_features_lst = pickle.load(f)

Human_HBB_nparray_dict = HBB_features_lst[0]
Human_HBB_tensor_dict = {k: tf.constant(v) for k, v in Human_HBB_nparray_dict.items()}
protein = Human_HBB_tensor_dict
print(f"特征：{protein.keys()}")
print(protein['msa'])
print(protein['template_aatype']) 

new_order_list = MAP_HHBLITS_AATYPE_TO_OUR_AATYPE
new_order = tf.constant(new_order_list, dtype=protein['msa'].dtype)
print(f"new_order:{new_order}")
print(f"new_order type :{type(new_order)}")
print(f"new_order shape :{(new_order.shape)}")

## 调用映射函数
fixed_protein = correct_msa_restypes(protein)
fixed_protein = fix_templates_aatype(protein)
print("fixed_protein")
print(fixed_protein['msa'])
print(fixed_protein['template_aatype']) 
#print(fixed_protein)