Given p0-p3, compute dihedral b/t planes p0p1p2 and p1p2p3.
(p0, p1, p2, p3)
| 632 | |
| 633 | |
| 634 | def get_dihedral(p0, p1, p2, p3): |
| 635 | """ |
| 636 | Given p0-p3, compute dihedral b/t planes p0p1p2 and p1p2p3. |
| 637 | """ |
| 638 | assert p0.shape[-1] == p1.shape[-1] == p2.shape[-1] == p3.shape[-1] == 3 |
| 639 | |
| 640 | # dx |
| 641 | b0 = p1 - p0 |
| 642 | b1 = p2 - p1 |
| 643 | b2 = p3 - p2 |
| 644 | |
| 645 | # normals |
| 646 | n012 = torch.cross(b0, b1) |
| 647 | n123 = torch.cross(b1, b2) |
| 648 | |
| 649 | # dihedral |
| 650 | cos_theta = torch.einsum('...i,...i->...', n012, n123) / ( |
| 651 | torch.norm(n012, dim=-1) * torch.norm(n123, dim=-1) + 1e-10) |
| 652 | sin_theta = torch.einsum('...i,...i->...', torch.cross(n012, n123), b1) / ( |
| 653 | torch.norm(n012, dim=-1) * torch.norm(n123, dim=-1) * torch.norm(b1, dim=-1) + 1e-10) |
| 654 | theta = torch.atan2(sin_theta, cos_theta) |
| 655 | |
| 656 | return theta |
| 657 | |
| 658 | |
| 659 | def get_chi_mask(protein, chi_id=None): |