(protein, rotate_angles)
| 732 | |
| 733 | @torch.no_grad() |
| 734 | def rotate_side_chain(protein, rotate_angles): |
| 735 | assert rotate_angles.shape[0] == protein.num_residue |
| 736 | assert rotate_angles.shape[1] == 4 |
| 737 | node_position14, mask14 = get_atom14_position(protein) # (num_residue, 14, 3) |
| 738 | |
| 739 | chi_atom14_index = chi_atom14_index_map.to(protein.device)[protein.residue_type] # (num_residue, 4, 4) 0~13 |
| 740 | chi_atom14_mask = chi_atom14_index != -1 |
| 741 | chi_atom14_index[~chi_atom14_mask] = 0 |
| 742 | for i in range(4): |
| 743 | atom_1, atom_2, atom_3, atom_4 = chi_atom14_index[:, i, :].unbind(-1) # (num_residue, ) |
| 744 | atom_2_position = torch.gather(node_position14, -2, |
| 745 | atom_2[:, None, None].expand(-1, -1, 3)) # (num_residue, 1, 3) |
| 746 | atom_3_position = torch.gather(node_position14, -2, |
| 747 | atom_3[:, None, None].expand(-1, -1, 3)) # (num_residue, 1, 3) |
| 748 | axis = atom_3_position - atom_2_position |
| 749 | axis_normalize = axis / (axis.norm(dim=-1, keepdim=True) + 1e-10) |
| 750 | rotate_angle = rotate_angles[:, i, None, None] |
| 751 | |
| 752 | # Rotate all subsequent atoms by the rotation angle |
| 753 | rotate_atoms_position = node_position14 - atom_2_position # (num_residue, 14, 3) |
| 754 | parallel_component = (rotate_atoms_position * axis_normalize).sum(dim=-1, keepdim=True) \ |
| 755 | * axis_normalize |
| 756 | perpendicular_component = rotate_atoms_position - parallel_component |
| 757 | perpendicular_component_norm = perpendicular_component.norm(dim=-1, keepdim=True) + 1e-10 |
| 758 | perpendicular_component_normalize = perpendicular_component / perpendicular_component_norm |
| 759 | normal_vector = torch.cross(axis_normalize.expand(-1, 14, -1), perpendicular_component_normalize, dim=-1) |
| 760 | transformed_atoms_position = perpendicular_component * rotate_angle.cos() + \ |
| 761 | normal_vector * perpendicular_component_norm * rotate_angle.sin() + \ |
| 762 | parallel_component + atom_2_position # (num_residue, 14, 3) |
| 763 | assert not transformed_atoms_position.isnan().any() |
| 764 | chi_mask = chi_atom14_mask[:, i, :].all(dim=-1, keepdim=True) # (num_residue, 1) |
| 765 | atom_mask = torch.arange(14, device=protein.device)[None, :] >= atom_4[:, None] # (num_residue, 14) |
| 766 | mask = (atom_mask & chi_mask).unsqueeze(-1).expand_as(node_position14) |
| 767 | node_position14[mask] = transformed_atoms_position[mask] |
| 768 | |
| 769 | protein.node_position[mask14] = node_position14[protein.atom2residue[mask14], protein.atom14index[mask14]] |
| 770 | return chi_atom14_mask.all(dim=-1) |
| 771 | |
| 772 | |
| 773 | @torch.no_grad() |
no test coverage detected