AF3 convert_to_af3_simple_features函数解读

AlphaFold3 data_transforms 模块 convert_to_af3_simple_features 函数的作用是提取 AlphaFold3 训练所需的简化骨架（backbone）特征，具体来说：

• 只保留 前 4 个骨架原子的坐标（N, CA, C, O）。

• 过滤掉侧链原子，使数据更简单，适用于 简化训练。

源代码：

def convert_to_af3_simple_features(protein):
    """Extracts the backbone features for AlphaFold3 simplified training.
    This is a temporary fix, will be removed later."""
    num_res = protein["atom14_atom_exists"].shape[0]

    # Add the atom_to_token mapping  (at this stage, cleaner to add this upstream)
    # res_to_token = torch.arange(num_res).unsqueeze(-1)  # (num_res, 1)
    # protein["atom_to_token"] = res_to_token.expand((-1, 4))  # (num_res, 4)

    # Crop the atom positions to the first 4 backbone atoms
    protein["all_atom_positions"] = protein["all_atom_positions"][:, :4, :]
    protein["all_atom_mask"] = protein["all_atom_mask"][:, :4]  # .reshape(num_res * 4)
    protein["all_atom_exists"] = protein["atom14_gt_exists"][:, :4]  # .reshape(num_res * 4)

    return protein

源码解读：

函数签名

def convert_to_af3_simple_features(protein):
    """Extracts the backbone features for AlphaFold3 simplified training.
    This is a temporary fix, will be removed later."""

• 函数名：convert_to_af3_simple_features

• 输入：protein —— 一个包含蛋白质特征信息的字典。

• 作用：转换 protein 特征数据，使其仅包含骨架（N, CA, C, O）的特征，去掉侧链信息。

计算残基数量

num_res = protein["atom14_atom_exists"].shape[0]

获取蛋白质的残基数：

• protein["atom14_atom_exists"] 是一个 (num_res, 14) 的数组。

• 其中 num_res 是氨基酸残基数量，14 代表 最多 14 个原子（包括 backbone 和 side chain）。

• 这里取 shape[0]，获取 残基数。

只保留骨架原子（N, CA, C, O）

protein["all_atom_positions"] = protein["all_atom_positions"][:, :4, :]
protein["all_atom_mask"] = protein["all_atom_mask"][:, :4]
protein["all_atom_exists"] = protein["atom14_gt_exists"][:, :4]

修改 protein 结构，使其只包含 Backbone（N, CA, C, O）：

• protein["all_atom_positions"]:

  • 形状 (num_res, 14, 3) → (num_res, 4, 3)

  • 只保留前 4 个原子（N, CA, C, O），丢弃侧链原子。

• protein["all_atom_mask"]:

  • 形状 (num_res, 14) → (num_res, 4)

  • 表示每个原子的存在情况（1 = 存在，0 = 缺失）。

• protein["all_atom_exists"]:

  • 形状 (num_res, 14) → (num_res, 4)

  • 也是一个 存在性 mask，和 atom_mask 类似，作用是标记哪些原子 实际存在（对于某些氨基酸，氧原子可能缺失）。

总结

🔹 该函数的作用

1. 简化蛋白质特征，去除侧链信息，只保留骨架原子（N, CA, C, O）。

2. 减少计算量，适用于轻量级训练或 Transformer 处理 Backbone-only 结构。

3. 修改 protein 数据结构，确保所有骨架特征 对齐。

🔹 在 AlphaFold3 训练中的作用

• 在 特定实验设置（如 Backbone-only 预测任务）中，只用 Backbone 信息训练模型。

• 适用于 简化的 AlphaFold3 训练 pipeline，可能用于快速验证 Backbone 结构预测的效果。