Returns a transformation object from reference coordinates. Note that this method does not take care of symmetries. If you provide the atom positions in the non-standard way, the N atom will end up not at [-0.527250, 1.359329, 0.0] but instead at [-0.527250, -1.359329, 0.0]. Yo
(n_xyz, ca_xyz, c_xyz, eps=1e-20)
| 858 | |
| 859 | @torch.no_grad() |
| 860 | def get_rigid_transform(n_xyz, ca_xyz, c_xyz, eps=1e-20): |
| 861 | """ |
| 862 | Returns a transformation object from reference coordinates. |
| 863 | |
| 864 | Note that this method does not take care of symmetries. If you |
| 865 | provide the atom positions in the non-standard way, the N atom will |
| 866 | end up not at [-0.527250, 1.359329, 0.0] but instead at |
| 867 | [-0.527250, -1.359329, 0.0]. You need to take care of such cases in |
| 868 | your code. |
| 869 | |
| 870 | Args: |
| 871 | n_xyz: A [*, 3] tensor of nitrogen xyz coordinates. |
| 872 | ca_xyz: A [*, 3] tensor of carbon alpha xyz coordinates. |
| 873 | c_xyz: A [*, 3] tensor of carbon xyz coordinates. |
| 874 | Returns: |
| 875 | A transformation (rots, translations). After applying the translation and |
| 876 | rotation to the reference backbone, the coordinates will |
| 877 | approximately equal to the input coordinates. |
| 878 | """ |
| 879 | |
| 880 | translation = -1 * ca_xyz |
| 881 | n_xyz = n_xyz + translation |
| 882 | c_xyz = c_xyz + translation |
| 883 | |
| 884 | c_x, c_y, c_z = [c_xyz[..., i] for i in range(3)] |
| 885 | norm = torch.sqrt(eps + c_x ** 2 + c_y ** 2) |
| 886 | sin_c1 = -c_y / norm |
| 887 | cos_c1 = c_x / norm |
| 888 | zeros = sin_c1.new_zeros(sin_c1.shape) |
| 889 | ones = sin_c1.new_ones(sin_c1.shape) |
| 890 | |
| 891 | c1_rots = sin_c1.new_zeros((*sin_c1.shape, 3, 3)) |
| 892 | c1_rots[..., 0, 0] = cos_c1 |
| 893 | c1_rots[..., 0, 1] = -1 * sin_c1 |
| 894 | c1_rots[..., 1, 0] = sin_c1 |
| 895 | c1_rots[..., 1, 1] = cos_c1 |
| 896 | c1_rots[..., 2, 2] = 1 |
| 897 | |
| 898 | norm = torch.sqrt(eps + c_x ** 2 + c_y ** 2 + c_z ** 2) |
| 899 | sin_c2 = c_z / norm |
| 900 | cos_c2 = torch.sqrt(c_x ** 2 + c_y ** 2) / norm |
| 901 | |
| 902 | c2_rots = sin_c2.new_zeros((*sin_c2.shape, 3, 3)) |
| 903 | c2_rots[..., 0, 0] = cos_c2 |
| 904 | c2_rots[..., 0, 2] = sin_c2 |
| 905 | c2_rots[..., 1, 1] = 1 |
| 906 | c2_rots[..., 2, 0] = -1 * sin_c2 |
| 907 | c2_rots[..., 2, 2] = cos_c2 |
| 908 | |
| 909 | c_rots = rot_matmul(c2_rots, c1_rots) |
| 910 | n_xyz = rot_vec_mul(c_rots, n_xyz) |
| 911 | |
| 912 | _, n_y, n_z = [n_xyz[..., i] for i in range(3)] |
| 913 | norm = torch.sqrt(eps + n_y ** 2 + n_z ** 2) |
| 914 | sin_n = -n_z / norm |
| 915 | cos_n = n_y / norm |
| 916 | |
| 917 | n_rots = sin_c2.new_zeros((*sin_c2.shape, 3, 3)) |
no test coverage detected