AF3 _make_rigid_group_constants函数解读

AlphaFold3 residue_constants 模块 _make_rigid_group_constants 函数负责 初始化氨基酸残基的刚体坐标系转换信息，具体包括：

1. 构造 restype_atom37_to_rigid_group 等映射关系

   • 确定氨基酸的 每个原子属于哪个刚性组（rigid group）

   • 记录原子的相对位置（在局部坐标系中的坐标）

2. 构造 restype_rigid_group_default_frame

   • 计算并存储从 backbone frame（主链框架）到 phi-frame、psi-frame、chi-frames 的刚体变换矩阵，用于不同 frame 之间的坐标转换。

源代码：

def _make_rigid_group_constants():
    """Fill the arrays above."""
    for restype, restype_letter in enumerate(restypes):
        resname = restype_1to3[restype_letter]
        for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]:
            atomtype = atom_order[atomname]
            restype_atom37_to_rigid_group[restype, atomtype] = group_idx
            restype_atom37_mask[restype, atomtype] = 1
            restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position

            atom14idx = restype_name_to_atom14_names[resname].index(atomname)
            restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
            restype_atom14_mask[restype, atom14idx] = 1
            restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position

    for restype, restype_letter in enumerate(restypes):
        resname = restype_1to3[restype_letter]
        atom_positions = {name: np.array(pos) for name, _, pos
                          in rigid_group_atom_positions[resname]}

        # backbone to backbone is the identity transform
        restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)

        # pre-omega-frame to backbone (currently dummy identity matrix)
        restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)

        # phi-frame to backbone
        mat = _make_rigid_transformation_4x4(
            ex=atom_positions['N'] - atom_positions['CA'],
            ey=np.array([1., 0., 0.]),
            translation=atom_positions['N'])
        restype_rigid_group_default_frame[restype, 2, :, :] = mat

        # psi-frame to backbone
        mat = _make_rigid_transformation_4x4(
            ex=atom_positions['C'] - atom_positions['CA'],
            ey=atom_positions['CA'] - atom_positions['N'],
            translation=atom_positions['C'])
        restype_rigid_group_default_frame[restype, 3, :, :] = mat

        # chi1-frame to backbone
        if chi_angles_mask[restype][0]:
            base_atom_names = chi_angles_atoms[resname][0]
            base_atom_positions = [atom_positions[name] for name in base_atom_names]
            mat = _make_rigid_transformation_4x4(
                ex=base_atom_positions[2] - base_atom_positions[1],
                ey=base_atom_positions[0] - base_atom_positions[1],
                translation=base_atom_positions[2])
            restype_rigid_group_default_frame[restype, 4, :, :] = mat

        # chi2-frame to chi1-frame
        # chi3-frame to chi2-frame
        # chi4-frame to chi3-frame
        # luckily all rotation axes for the next frame start at (0,0,0) of the
        # previous frame
        for chi_idx in range(1, 4):
            if chi_angles_mask[restype][chi_idx]:
                axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
                axis_end_atom_position = atom_positions[axis_end_atom_name]
                mat = _make_rigid_transformation_4x4(
                    ex=axis_end_atom_position,
                    ey=np.array([-1., 0., 0.]),
                    translation=axis_end_atom_position)
                restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat

源码解读：

1. 解析刚性组信息

for restype, restype_letter in enumerate(restypes):
    resname = restype_1to3[restype_letter]
    for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]:
        atomtype = atom_order[atomname]
        restype_atom37_to_rigid_group[restype, atomtype] = group_idx
        restype_atom37_mask[restype, atomtype] = 1
        restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position

        atom14idx = restype_name_to_atom14_names[resname].index(atomname)
        restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
        restype_atom14_mask[restype, atom14idx] = 1
        restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position

这部分代码的功能

• 通过 rigid_group_atom_positions 获取 每个氨基酸中每个原子的局部坐标

• 记录 每个原子属于哪个刚性组（rigid group）

• 存储 该原子在局部坐标系（frame）中的坐标

• 处理 Atom37（37 种可能的原子）和 Atom14（AlphaFold 只使用的 14 种原子）两种不同的表示方式

📌 为什么要做这些？

在 AlphaFold3 中，每个氨基酸的不同原子并不是完全独立的，而是按照 刚性组 来组织的。例如：

• 主链（backbone）原子（N, CA, C, O）形成一个刚性组，不能相对移动

• 侧链（sidechain）原子 则根据不同氨基酸的化学结构，分成不同的 chi-angle frame

• 通过 刚性组 的划分，可以用 一组旋转角度（如 phi, psi, chi）来控制整个结构，从而减少计算复杂度

2. 计算并存储默认的刚体变换矩阵

for restype, restype_letter in enumerate(restypes):
    resname = restype_1to3[restype_letter]
    atom_positions = {name: np.array(pos) for name, _, pos
                      in rigid_group_atom_positions[resname]}

这里先获取每个氨基酸的原子局部坐标，以便后续计算变换矩阵。

2.1 Backbone Frame（主链框架&#x