ProteinMPNN数据输入主要函数介绍

ProteinMPNN 是一种深度学习模型，专为蛋白质设计和序列预测任务开发。其核心思想是通过学习蛋白质结构和序列之间的映射关系，来生成新蛋白质序列或优化现有蛋白质的序列。ProteinMPNN 使用图神经网络（Graph Neural Networks, GNNs）来建模蛋白质的三维结构，并通过这些结构信息来预测每个位置上可能的氨基酸序列。

1. PDB_dataset类

PDB_dataset为自定义数据加载类（继承torch.utils.data.Dataset），先随机选择同源结构中的一个，self.train_dict[ID]为同源结构的列表如 ["8p0n_H","8p0j_H"]，再调用loader_pdb函数加载数据。通过torch.utils.data.DataLoader加载PDB_dataset。

class PDB_dataset(torch.utils.data.Dataset):
    def __init__(self, IDs, loader, train_dict, params):
        self.IDs = IDs
        self.train_dict = train_dict
        self.loader = loader
        self.params = params

    def __len__(self):
        return len(self.IDs)

    def __getitem__(self, index):
        ID = self.IDs[index]
        sel_idx = np.random.randint(0, len(self.train_dict[ID]))
        out = self.loader(self.train_dict[ID][sel_idx], self.params)
        return out

2. loader_pdb函数

loader_pdb函数根据PDB ID号以及链的编号读入处理好的结构数据（ .pt文件，关于这个结构的meta字典数据以及各个链chain的数据），随机选择一个组装体，进行坐标变换，根据同源性对链进行掩码mask设计 ，最后合并一个结构中多条链的序列以及坐标的合并，输出字典，包含seq' , 'xyz' , 'idx',  'masked'，'label' 键。

def loader_pdb(item,params):

    pdbid,chid = item[0].split('_')
    PREFIX = "%s/pdb/%s/%s"%(params['DIR'],pdbid[1:3],pdbid)
    
    # load metadata
    if not os.path.isfile(PREFIX+".pt"):
        return {'seq': np.zeros(5)}
    meta = torch.load(PREFIX+".pt")
    asmb_ids = meta['asmb_ids']
    asmb_chains = meta['asmb_chains']
    chids = np.array(meta['chains'])

    # find candidate assemblies which contain chid chain
    asmb_candidates = set([a for a,b in zip(asmb_ids,asmb_chains)
                           if chid in b.split(',')])

    # if the chains is missing is missing from all the assemblies
    # then return this chain alone
    if len(asmb_candidates)<1:
        chain = torch.load("%s_%s.pt"%(PREFIX,chid))
        L = len(chain['seq'])
        return {'seq'    : chain['seq'],
                'xyz'    : chain['xyz'],
                'idx'    : torch.zeros(L).int(),
                'masked' : torch.Tensor([0]).int(),
                'label'  : item[0]}

    # randomly pick one assembly from candidates
    asmb_i = random.sample(list(asm