AlphaFold3 protein_datamodule 模块 Reorder
类的作用确实是将蛋白质骨架原子的顺序从一种格式转换成另一种格式,将 ProteinFlow
来源的数据格式标准化为 AlphaFold3 所期望的顺序。
源代码:
class Reorder(nn.Module):"""A transformation that reorders the 3D coordinates of backbone atomsfrom N, C, Ca, O -> N, Ca, C, O."""def forward(self, protein_dict):# Switch to N, Ca, C, ordering.reordered_X = protein_dict['X'].index_select(1, torch.tensor([0, 2, 1, 3]))protein_dict['X'] = reordered_Xreturn protein_dict
源码解读:
类定义和功能
class Reorder(nn.Module):"""A transformation that reorders the 3D coordinates of backbone atomsfrom N, C, Ca, O -> N, Ca, C, O."""
-
这是一个继承自
torch.nn.Module
的可调用模块,用于变换蛋白质原子顺序。 -
文档字符串中已经明确指出它做的事:从
[N, C, Ca, O]
顺序 →[N, Ca, C, O]
顺序。
forward 方法解释
def forward(self, protein_dict):# Switch to N, Ca, C, ordering.reordered_X = protein_dict['X'].index_select(1, torch.tensor([0, 2, 1, 3]))protein_dict['X'] = reordered_Xreturn protein_dict
-
protein_dict
:一个包含蛋白质结构信息的字典。字段X
存储原子坐标张量。 -
protein_dict['X']
的形状通常是(L, 4, 3)—— 单个样本,不含批次维度
-
index_select(1, torch.tensor([0, 2, 1, 3]))
的作用是重新排列第2个维度(即原子顺序),具体来说:-
第0维是残基(或批次),第1维是原子,索引为 [N, C, Ca, O]
-
新顺序
[0, 2, 1, 3]
就是 →[N, Ca, C, O]
-
-
注:torch.index_select
是 PyTorch 中一个非常实用的张量操作方法,用于按指定的索引在某一维度上选取元素。这个操作在数据处理、切片、序列重排等任务中非常常见。
ProteinFlow中原子
顺序:
[N, C, Ca, O] # 也就是 [0, 1, 2, 3]
AlphaFold3 所期望的顺序:
[N, Ca, C, O] # 也就是 [0, 2, 1, 3]
最后,将修改后的 X
重新赋值回 protein_dict['X']
并返回。