Initialize the side chain of the protein from restype_atom14_rigid_group_positions (shape: 21, 14, 3) Args: protein: Protein object keep_backbone: If True, the backbone will not be changed
(protein, keep_backbone=True)
| 931 | |
| 932 | @torch.no_grad() |
| 933 | def init_sidechain(protein, keep_backbone=True): |
| 934 | """ |
| 935 | Initialize the side chain of the protein from restype_atom14_rigid_group_positions (shape: 21, 14, 3) |
| 936 | Args: |
| 937 | protein: Protein object |
| 938 | keep_backbone: If True, the backbone will not be changed |
| 939 | """ |
| 940 | node_position14, mask14 = get_atom14_position(protein) # (num_residue, 14, 3), (num_residue, 14) |
| 941 | assert mask14[..., 0].all() # Each Residue should have N atom |
| 942 | assert mask14[..., 1].all() # Each Residue should have CA atom |
| 943 | assert mask14[..., 2].all() # Each Residue should have C atom |
| 944 | rots, translation = get_rigid_transform(n_xyz=node_position14[..., 0, :], |
| 945 | ca_xyz=node_position14[..., 1, :], |
| 946 | c_xyz=node_position14[..., 2, :]) # (num_residue, 3, 3), (num_residue, 3) |
| 947 | # Align the rigid side chain to backbone |
| 948 | new_node_position14 = rot_vec_mul(rots.unsqueeze(-3), restype_atom14_rigid_group_positions[protein.residue_type]) \ |
| 949 | + translation.unsqueeze(-2) |
| 950 | new_node_position14 = new_node_position14 * restype_atom14_mask[protein.residue_type].unsqueeze(-1) |
| 951 | |
| 952 | # Preserve the original position of backbone atoms |
| 953 | if keep_backbone: |
| 954 | new_node_position14[..., 0:4] = node_position14[..., 0:4] |
| 955 | |
| 956 | # Set the position of 14 atoms for each residue |
| 957 | protein = set_atom14_position(protein, new_node_position14) |
| 958 | return protein |
| 959 | |
| 960 | |
| 961 | @torch.no_grad() |
nothing calls this directly
no test coverage detected