函数功能
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: 布尔掩码，指示哪些扭转角是平面的。
源代码：
def get_torsions(xyz_in, seq, torsion_indices, torsion_can_flip, ref_angles, mask_in=None
):B, L = xyz_in.shape[:2]tors_mask = get_tor_mask(seq, torsion_indices, mask_in)# torsions to restrain to 0 or 180degreetors_planar = torch.zeros((B, L, 10), dtype=torch.bool, device=xyz_in.device)tors_planar[:, :, 5] = seq == aa2num["TYR"] # TYR chi 3 should be planar# idealize given xyz coordinates before computing torsion anglesxyz = xyz_in.clone()Rs, Ts = rigid_from_3_points(xyz[..., 0, :], xyz[..., 1, :], xyz[..., 2, :])Nideal = torch.tensor([-0.5272, 1.3593, 0.000], device=xyz_in.device)Cideal = torch.tensor([1.5233, 0.000, 0.000], device=xyz_in.device)xyz[..., 0, :] = torch.einsum("brij,j->bri", Rs, Nideal) + Tsxyz[..., 2, :] = torch.einsum("brij,j->bri", Rs, Cideal) + Tstorsions = torch.zeros((B, L, 10, 2), device=xyz.device)# avoid undefined angles for H generationtorsions[:, 0, 1, 0] = 1.0torsions[:, -1, 0, 0] = 1.0# omegatorsions[:, :-1, 0, :] = th_dih(xyz[:, :-1, 1, :], xyz[:, :-1, 2, :], xyz[:, 1:, 0, :], xyz[:, 1:, 1, :])# phitorsions[:, 1:, 1, :] = th_dih(xyz[:, :-1, 2, :], xyz[:, 1:, 0, :], xyz[:, 1:, 1, :], xyz[:, 1:, 2, :])# psitorsions[:, :, 2, :] = -1 * th_dih(xyz[:, :, 0, :], xyz[:, :, 1, :], xyz[:, :, 2, :], xyz[:, :, 3, :])# chisti0 = torch.gather(xyz, 2, torsion_indices[seq, :, 0, None].repeat(1, 1, 1, 3))ti1 = torch.gather(xyz, 2, torsion_indices[seq, :, 1, None].repeat(1, 1, 1, 3))ti2 = torch.gather(xyz, 2, torsion_indices[seq, :, 2, None].repeat(1, 1, 1, 3))ti3 = torch.gather(xyz, 2, torsion_indices[seq, :, 3, None].repeat(1, 1, 1, 3))torsions[:, :, 3:7, :] = th_dih(ti0, ti1, ti2, ti3)# CB bendNC = 0.5 * (xyz[:, :, 0, :3] + xyz[:, :, 2, :3])CA = xyz[:, :, 1, :3]CB = xyz[:, :, 4, :3]t = th_ang_v(CB - CA, NC - CA)t0 = ref_angles[seq][..., 0, :]torsions[:, :, 7, :] = torch.stack((torch.sum(t * t0, dim=-1), t[..., 0] * t0[..., 1] - t[..., 1] * t0[..., 0]),dim=-1,)# CB twistNCCA = NC - CANCp = xyz[:, :, 2, :3] - xyz[:, :, 0, :3]NCpp = (NCp- torch.sum(NCp * NCCA, dim=-1, keepdim=True)/ torch.sum(NCCA * NCCA, dim=-1, keepdim=True)* NCCA)t = th_ang_v(CB - CA, NCpp)t0 = ref_angles[seq][..., 1, :]torsions[:, :, 8, :] = torch.stack((torch.sum(t * t0, dim=-1), t[..., 0] * t0[..., 1] - t[..., 1] * t0[..., 0]),dim=-1,)# CG bendCG = xyz[:, :
集成)
