RFdiffusion get_torsions函数解读

函数功能
get_torsions 函数根据输入的原子坐标（xyz_in）和氨基酸序列（seq），盘算一组主链和侧链的扭转角（torsions）。同时天生备用扭转角（torsions_alt），用于表现可以镜像翻转的多少结构，并返回掩码（tors_mask）和是否平面化（tors_planar）的信息。
输入参数

  

• xyz_in: 原子坐标张量，外形 [B, L, 27, 3]，此中 B 是批量巨细，L 是残基数目，27 表现原子范例（如 N, CA, C, CB），3 是坐标。  
  
• seq: 氨基酸序列，外形 [B, L]，每个值对应氨基酸编号。  
  
• torsion_indices: 表现盘算侧链扭转角所需的 4 个原子索引。  
  
• torsion_can_flip: 布尔数组，指示哪些扭转角可以翻转。  
  
• ref_angles: 抱负化的参考角度，重要用于 CB 弯曲、CB 扭转和 CG 弯曲。  
  
• mask_in: 掩码（可选），用于屏蔽特定残基。
  

返回值

  

• torsions: 主链和侧链的扭转角张量，外形 [B, L, 10, 2]。  
  
• torsions_alt: 扭转角备用版本（翻转版），外形同上。  
  
• tors_mask: 扭转角的有用掩码。  
  
• tors_planar: 布尔掩码，指示哪些扭转角是平面的。
  

源代码：

1. def get_torsions(
2.     xyz_in, seq, torsion_indices, torsion_can_flip, ref_angles, mask_in=None
3. ):
4.     B, L = xyz_in.shape[:2]
6.     tors_mask = get_tor_mask(seq, torsion_indices, mask_in)
8.     # torsions to restrain to 0 or 180degree
9.     tors_planar = torch.zeros((B, L, 10), dtype=torch.bool, device=xyz_in.device)
10.     tors_planar[:, :, 5] = seq == aa2num["TYR"]  # TYR chi 3 should be planar
12.     # idealize given xyz coordinates before computing torsion angles
13.     xyz = xyz_in.clone()
14.     Rs, Ts = rigid_from_3_points(xyz[..., 0, :], xyz[..., 1, :], xyz[..., 2, :])
15.     Nideal = torch.tensor([-0.5272, 1.3593, 0.000], device=xyz_in.device)
16.     Cideal = torch.tensor([1.5233, 0.000, 0.000], device=xyz_in.device)
17.     xyz[..., 0, :] = torch.einsum("brij,j->bri", Rs, Nideal) + Ts
18.     xyz[..., 2, :] = torch.einsum("brij,j->bri", Rs, Cideal) + Ts
20.     torsions = torch.zeros((B, L, 10, 2), device=xyz.device)
21.     # avoid undefined angles for H generation
22.     torsions[:, 0, 1, 0] = 1.0
23.     torsions[:, -1, 0, 0] = 1.0
25.     # omega
26.     torsions[:, :-1, 0, :] = th_dih(
27.         xyz[:, :-1, 1, :], xyz[:, :-1, 2, :], xyz[:, 1:, 0, :], xyz[:, 1:, 1, :]
28.     )
29.     # phi
30.     torsions[:, 1:, 1, :] = th_dih(
31.         xyz[:, :-1, 2, :], xyz[:, 1:, 0, :], xyz[:, 1:, 1, :], xyz[:, 1:, 2, :]
32.     )
33.     # psi
34.     torsions[:, :, 2, :] = -1 * th_dih(
35.         xyz[:, :, 0, :], xyz[:, :, 1, :], xyz[:, :, 2, :], xyz[:, :, 3, :]
36.     )
38.     # chis
39.     ti0 = torch.gather(xyz, 2, torsion_indices[seq, :, 0, None].repeat(1, 1, 1, 3))
40.     ti1 = torch.gather(xyz, 2, torsion_indices[seq, :, 1, None].repeat(1, 1, 1, 3))
41.     ti2 = torch.gather(xyz, 2, torsion_indices[seq, :, 2, None].repeat(1, 1, 1, 3))
42.     ti3 = torch.gather(xyz, 2, torsion_indices[seq, :, 3, None].repeat(1, 1, 1, 3))
43.     torsions[:, :, 3:7, :] = th_dih(ti0, ti1, ti2, ti3)
45.     # CB bend
46.     NC = 0.5 * (xyz[:, :, 0, :3] + xyz[:, :, 2, :3])
47.     CA = xyz[:, :, 1, :3]
48.     CB = xyz[:, :, 4, :3]
49.     t = th_ang_v(CB - CA, NC - CA)
50.     t0 = ref_angles[seq][..., 0, :]
51.     torsions[:, :, 7, :] = torch.stack(
52.         (torch.sum(t * t0, dim=-1), t[..., 0] * t0[..., 1] - t[..., 1] * t0[..., 0]),
53.         dim=-1,
54.     )
56.     # CB twist
57.     NCCA = NC - CA
58.     NCp = xyz[:, :, 2, :3] - xyz[:, :, 0, :3]
59.     NCpp = (
60.         NCp
61.         - torch.sum(NCp * NCCA, dim=-1, keepdim=True)
62.         / torch.sum(NCCA * NCCA, dim=-1, keepdim=True)
63.         * NCCA
64.     )
65.     t = th_ang_v(CB - CA, NCpp)
66.     t0 = ref_angles[seq][..., 1, :]
67.     torsions[:, :, 8, :] = torch.stack(
68.         (torch.sum(t * t0, dim=-1), t[..., 0] * t0[..., 1] - t[..., 1] * t0[..., 0]),
69.         dim=-1,
70.     )
72.     # CG bend
73.     CG = xyz[:, :, 5, :3]
74.     t = th_ang_v(CG - CB, CA - CB)
75.     t0 = ref_angles[seq][..., 2, :]
76.     torsions[:, :, 9, :] = torch.stack(
77.    

复制代码