函數功能
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[:, :
集成)
