File size: 1,709 Bytes
8e04e6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import torch
def normalize_last_dim(v):
norm = torch.linalg.norm(v, dim=-1, keepdim=True)
return v / torch.clamp(norm, min=1e-8, max=None)
def bond_angles(a, b, c):
"""
Computes bond angles for the 3 points a, b, c.
Since torch.linalg.cross and torch.linalg.cross support
broadcasting, this supports broadcasting.
Args:
a, b, c: Each is a tensor of shape [*, 3]
Returns:
Angle between 0 and pi, shape [*]
"""
b0 = b - a # [*, 3]
b1 = c - a # [*, 3]
b0, b1 = map(normalize_last_dim, (b0, b1)) # [*, 3] each
cos_angle = torch.linalg.vecdot(b0, b1, dim=-1) # [*]
cross = torch.linalg.cross(b0, b1, dim=-1) # [*, 3]
sin_angle = torch.linalg.norm(cross, dim=-1) # [*]
return torch.atan2(sin_angle, cos_angle) # [*]
def signed_dihedral_angle(a, b, c, d):
"""
Compputes the signed angle for the 4 points a, b, c, d.
Since torch.linalg.cross and torch.linalg.cross support
broadcasting, this supports broadcasting.
Args:
a, b, c, d: Each is a tensor of shape [*, 3]
Returns:
Angle between -pi and pi (signed), shape [*]
"""
b0 = b - a # [*, 3]
b1 = c - b # [*, 3]
b2 = d - c # [*, 3]
n1 = torch.linalg.cross(b0, b1) # [*, 3]
n2 = torch.linalg.cross(b1, b2) # [*, 3]
n1, n2 = map(normalize_last_dim, (n1, n2)) # Each [*, 3]
cos_angle = torch.linalg.vecdot(n1, n2, dim=-1) # [*]
n1_cross_n2 = torch.linalg.cross(n1, n2, dim=-1)
sin_angle_magnitude = torch.linalg.norm(n1_cross_n2, dim=-1) # [*]
sign = torch.sign(torch.linalg.vecdot(n1_cross_n2, b1)) # [*]
return torch.atan2(sign * sin_angle_magnitude, cos_angle) # [*]
|