RFDiffusion 计算键角函数get_ang解读

get_ang 函数（kinematics.py包中）计算三组原子 a,b,c 所形成的平面角（planar angle），即 b 为顶点， a,b,c 所确定的角度。

源代码：

def get_ang(a, b, c):
    """calculate planar angles for all consecutive triples (a[i],b[i],c[i])
    from Cartesian coordinates of three sets of atoms a,b,c 

    Parameters
    ----------
    a,b,c : pytorch tensors of shape [batch,nres,3]
            store Cartesian coordinates of three sets of atoms
    Returns
    -------
    ang : pytorch tensor of shape [batch,nres]
          stores resulting planar angles
    """
    v = a - b
    w = c - b
    v /= torch.norm(v, dim=-1, keepdim=True)
    w /= torch.norm(w, dim=-1, keepdim=True)
    vw = torch.sum(v*w, dim=-1)

    return torch.acos(vw)

c6d[b,i,j,torch.full_like(b,3)] = get_ang(Ca[b,i], Cb[b,i], Cb[b,j])

代码解读：

输入参数

• a, b, c：形状为 [batch, nres, 3] 的 PyTorch 张量，表示多个样本中各残基的三组原子笛卡尔坐标。
  • batch：批量大小。
  • nres：每个样本的残基数量。
  • 3：表示每个原子的三维坐标 (x,y,z)(x,y,z)。

输出

• ang：形状为 [batch, nres] 的张量，表示每个样本中残基对 a,b,ca,b,c 形成的平面角（弧度）。

步骤解析

1. 计算向量

   • v = a - b：从 bb 指向 aa 的向量。
   • w = c - b：从 bb 指向 cc 的向量。

2. 单位化向量

   • v /= torch.norm(v, dim=-1, keepdim=True)：将向量 vv 归一化。
   • w /= torch.norm(w, dim=-1, keepdim=True)：将向量 ww 归一化。

3. 计算夹角的余弦值

   • vw = torch.sum(v*w, dim=-1)：向量 vv 和 ww 的点积，结果是每个残基间的余弦值。

4. 余弦值转角度

   • torch.acos(vw)：使用反余弦函数计算角度，返回弧度制的角度。

          函数调用

c6d[b, i, j, torch.full_like(b, 3)] = get_ang(Ca[b, i], Cb[b, i], Cb[b, j])

• Ca[b, i]：第 bb 个样本中，第 i 个残基的 CαCα​ 坐标。
• Cb[b, i] 和 Cb[b, j]：第 b 个样本中，第 i 和第 j 个残基的 Cβ 坐标。
• get_ang(Ca[b, i], Cb[b, i], Cb[b, j])：计算第 b 个样本中，以 i 残基 Cβ 为顶点、以 i 残基 Cα 和 j 残基 Cβ 确定的角度。
• c6d[b, i, j, torch.full_like(b, 3)]：将计算结果存储到一个 6D 张量（可能用于结构表征）。