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