AF3 make_pseudo_beta函数解读

AlphaFold3  data_transforms 模块的 make_pseudo_beta 函数 通过调用函数pseudo_beta_fn 为蛋白质结构生成伪β碳（pseudo-beta，简称 pβ）的位置和掩码。伪β碳是一种用于表示蛋白质主链结构的简化原子，通常用于替代β碳（CB），而在甘氨酸（Glycine，简称 Gly）中则使用α碳（CA）作为替代。

源代码：

@curry1
def make_pseudo_beta(protein, prefix=""):
    """Create pseudo-beta (alpha for glycine) position and mask."""
    assert prefix in ["", "template_"]
    (
        protein[prefix + "pseudo_beta"],
        protein[prefix + "pseudo_beta_mask"],
    ) = pseudo_beta_fn(
        protein["template_aatype" if prefix else "aatype"],
        protein[prefix + "all_atom_positions"],
        protein["template_all_atom_mask" if prefix else "all_atom_mask"],
    )
    return protein


def pseudo_beta_fn(aatype, all_atom_positions, all_atom_mask):
    """Create pseudo beta features."""
    is_gly = torch.eq(aatype, rc.restype_order["G"])
    ca_idx = rc.atom_order["CA"]
    cb_idx = rc.atom_order["CB"]
    pseudo_beta = torch.where(
        torch.tile(is_gly[..., None], [1] * len(is_gly.shape) + [3]),
        all_atom_positions[...,