蛋白质序列特征字典

将蛋白质的序列等特征转化为特征字典形式：键的类型是字符串 (str)，值的类型是 NumPy 数组。

from typing import MutableMapping
import numpy as np

# 定义映射类型，存储特征（特征名称：numpy array）
# MutableMapping[str, np.ndarray]：键的类型是字符串 (str)，值的类型是 NumPy 数组 (np.ndarray)
FeatureDict = MutableMapping[str, np.ndarray]

# 组成蛋白质的氨基酸
restypes = [
    'A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P',
    'S', 'T', 'W', 'Y', 'V'
]

# X:unkown
restypes_with_x = restypes + ['X'] 

restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)}
# print(restype_order_with_x)

# 热独（onehot）编码函数
def sequence_to_onehot(
      sequence: str,
      mapping: Mapping[str, int],
      map_unknown_to_x: bool = False) -> np.ndarray:
    """Maps the given sequence into a one-hot encoded matrix.

    Args:
        sequence: An amino acid sequence.
        mapping: A dictionary mapping amino acids to integers.
        map_unknown_to_x: If True, any amino acid that is not in the mapping will be
          mapped to the unknown amino acid 'X'. If the mapping doesn't contain
          amino acid 'X', an error will be thrown. If False, any amino acid not in
          the mapping will throw an error.

    Returns:
        A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of
          the sequence.

    Raises:
        ValueError: If the mapping doesn't contain values from 0 to
          num_unique_aas - 1 without any gaps.
    """
    num_entries = max(mapping.values()) + 1

    if sorted(set(mapping.values())) != list(range(num_entries)):
      raise ValueError('The mapping must have values from 0 to num_unique_aas-1 '
                       'without any gaps. Got: %s' % sorted(mapping.values()))

    one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)

    for aa_index, aa_type in enumerate(sequence):
        if map_unknown_to_x:
            if aa_type.isalpha() and aa_type.isupper():
                aa_id = mapping.get(aa_type, mapping['X'])
            else:
                raise ValueError(f'Invalid character in the sequence: {aa_type}')
        else:
            aa_id = mapping[aa_type]
        one_hot_arr[aa_index, aa_id] = 1

    return one_hot_arr


def make_sequence_features(
      sequence: str, description: str, num_res: int) -> FeatureDict:
    """Constructs a feature dict of sequence features."""
    features = {}
    features['aatype'] = sequence_to_onehot(
      sequence=sequence,
      mapping=restype_order_with_x,
      map_unknown_to_x=True)
    features['between_segment_residues'] = np.zeros((num_res,), dtype=np.int32)
    features['domain_name'] = np.array([description.encode('utf-8')],
                                        dtype=np.object_)
    features['residue_index'] = np.array(range(num_res), dtype=np.int32)
    features['seq_length'] = np.array([num_res] * num_res, dtype=np.int32)
    features['sequence'] = np.array([sequence.encode('utf-8')], dtype=np.object_)
    return features

seq = "ARNDCQEGHILKMFPSTWYVX"

seq_features = make_sequence_features(sequence = seq,
                                      description = "test",
                                      num_res =len(seq))

print(seq_features)