add the necessary functions for deriving the frames of the ligands, or other biomolecules that are not amino acids or nucleotides

Author: lucidrains

alphafold3_pytorch/inputs.py

@@ -139,6 +139,9 @@
 
 # functions
 
+def l2norm(t):
+    return F.normalize(t, dim = -1)
+
 def flatten(arr):
     return [el for sub_arr in arr for el in sub_arr]
 
@@ -199,6 +202,96 @@ def hard_validate_atom_indices_ascending(
 
         assert (difference >= 0).all(), f'detected invalid {error_msg_field} for in a batch: {present_indices}'
 
+# functions for deriving the frames for ligands
+# this follows the logic from Alphafold3 Supplementary section 4.3.2
+
+@typecheck
+def get_indices_three_closest_atom_pos(
+    atom_pos: Float['... n d'],
+) -> Int['... 3']:
+
+    prec_dims, device = atom_pos.shape[:-2], atom_pos.device
+    num_atoms, has_batch = atom_pos.shape[-2], atom_pos.ndim == 3
+
+    if num_atoms < 3:
+        return atom_pos.new_full((*prec_dims, 3), -1).long()
+
+    if not has_batch:
+        atom_pos = rearrange(atom_pos, '... -> 1 ...')
+
+    atom_dist = torch.cdist(atom_pos, atom_pos)
+
+    # mask out the distance to self
+
+    eye = torch.eye(num_atoms, device = device, dtype = torch.bool)
+
+    atom_dist.masked_fill_(eye, 1e4)
+
+    # will use topk on the negative of the distance
+
+    neg_distance, two_closest_atom_indices = (-atom_dist).topk(2, dim = -1)
+
+    mean_neg_distance = neg_distance.mean(dim = -1)
+
+    best_atom_pair_index = mean_neg_distance.argmax(dim = -1)
+
+    best_two_atom_neighbors = einx.get_at('... [m] c, ... -> ... c', two_closest_atom_indices, best_atom_pair_index)
+
+    # place the chosen atom at the center
+
+    three_atom_indices, _ = pack((
+        best_two_atom_neighbors[..., 0],
+        best_atom_pair_index,
+        best_two_atom_neighbors[..., 1],
+    ), 'b *')
+
+    if not has_batch:
+        three_atom_indices = rearrange(three_atom_indices, '1 ... -> ...')
+
+    return three_atom_indices
+
+@typecheck
+def get_angle_between_edges(
+    edge1: Float['... 3'],
+    edge2: Float['... 3']
+) -> Float['...']:
+    cos = torch.dot(l2norm(edge1), l2norm(edge2))
+    return torch.acos(cos)
+
+@typecheck
+def get_frames_from_atom_pos(
+    atom_pos: Float['... n d'],
+    filter_colinear_pos: bool = False,
+    is_colinear_angle_thres: float = 25. # they use 25 degrees as a way of filtering out invalid frames
+) -> Int['... 3']:
+
+    frames = get_indices_three_closest_atom_pos(atom_pos)
+
+    if not filter_colinear_pos:
+        return frames
+
+    # get the edges and derive angles
+
+    three_atom_pos = einx.get_at('... [m] c, ... three -> ... three c', atom_pos, frames)
+
+    left_pos, center_pos, right_pos = three_atom_pos.unbind(dim = -2)
+
+    edges1, edges2 = (left_pos - center_pos), (right_pos - center_pos)
+
+    angle = get_angle_between_edges(edges1, edges2)
+
+    degree = torch.rad2deg(angle)
+
+    is_colinear = (
+        (degree.abs() < is_colinear_angle_thres) |
+        ((180. - degree.abs()) < is_colinear_angle_thres)
+    )
+
+    # set any three atoms that are colinear to -1 indices
+
+    three_atom_indices = einx.where('..., ... three, -> ... three', ~is_colinear, frames, -1)
+    return three_atom_indices
+
 # atom level, what Alphafold3 accepts
 
 UNCOLLATABLE_ATOM_INPUT_FIELDS = {'filepath'}

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.4.14"
+version = "0.4.15"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" },

tests/test_af3.py

@@ -66,6 +66,9 @@
     PDBDataset,
     default_extract_atom_feats_fn,
     default_extract_atompair_feats_fn,
+    get_indices_three_closest_atom_pos,
+    get_angle_between_edges,
+    get_frames_from_atom_pos
 )
 
 from alphafold3_pytorch.utils.model_utils import exclusive_cumsum
@@ -192,6 +195,23 @@ def test_rigid_from_three_points():
     rotation, _ = rigid_from_3_points((points, points, points))
     assert rotation.shape == (7, 11, 23, 3, 3)
 
+def test_deriving_frames_for_ligands():
+    points = torch.tensor([
+        [1., 1., 1.],
+        [-99, -99, -99],
+        [0, 0, 0],
+        [100, 100, 100],
+        [-1., -1., -1.],
+    ])
+
+    frames = get_frames_from_atom_pos(points, filter_colinear_pos = True)
+
+    assert torch.allclose(frames, torch.tensor([-1, -1, -1]))
+
+    frames = get_frames_from_atom_pos(points, filter_colinear_pos = False)
+
+    assert torch.allclose(frames, torch.tensor([0, 2, 4]))
+
 def test_compute_alignment_error():
     pred_coords = torch.randn(2, 100, 3)
     pred_frames = torch.randn(2, 100, 3, 3)