distance_labels auto derived from atom positions + residue atom indices, if the latter two are passed in

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -2452,9 +2452,9 @@ def __init__(
         dim_single = 384,
         dim_pairwise = 128,
         dim_token = 768,
-        atompair_dist_bins: Float[' dist_bins'] = torch.linspace(3, 20, 37),
+        distance_bins: Float[' dist_bins'] = torch.linspace(3, 20, 38),
         ignore_index = -1,
-        num_dist_bins = 38,
+        num_dist_bins: int | None = None,
         num_plddt_bins = 50,
         num_pde_bins = 64,
         num_pae_bins = 64,
@@ -2626,14 +2626,19 @@ def __init__(
 
         # logit heads
 
+        self.register_buffer('distance_bins', distance_bins)
+        num_dist_bins = default(num_dist_bins, len(distance_bins))
+
+        assert len(distance_bins) == num_dist_bins, '`distance_bins` must have a length equal to the `num_dist_bins` passed in'
+
         self.distogram_head = DistogramHead(
             dim_pairwise = dim_pairwise,
             num_dist_bins = num_dist_bins
         )
 
         self.confidence_head = ConfidenceHead(
             dim_single_inputs = dim_single_inputs,
-            atompair_dist_bins = atompair_dist_bins,
+            atompair_dist_bins = distance_bins,
             dim_single = dim_single,
             dim_pairwise = dim_pairwise,
             num_plddt_bins = num_plddt_bins,
@@ -2830,6 +2835,12 @@ def forward(
 
         # distogram head
 
+        if not exists(distance_labels) and atom_pos_given and exists(residue_atom_indices):
+            residue_pos = einx.get_at('b (n [w]) c, b n -> b n c', atom_pos, residue_atom_indices)
+            residue_dist = torch.cdist(residue_pos, residue_pos, p = 2)
+            dist_from_dist_bins = einx.subtract('b m dist, dist_bins -> b m dist dist_bins', residue_dist, self.distance_bins).abs()
+            distance_labels = dist_from_dist_bins.argmin(dim = -1)
+
         if exists(distance_labels):
             distance_labels = torch.where(pairwise_mask, distance_labels, ignore)
             distogram_logits = self.distogram_head(pairwise)

pyproject.toml

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

tests/test_af3.py

@@ -384,7 +384,6 @@ def test_alphafold3():
     atom_pos = torch.randn(2, atom_seq_len, 3)
     residue_atom_indices = torch.randint(0, 27, (2, seq_len))
 
-    distance_labels = torch.randint(0, 38, (2, seq_len, seq_len))
     pae_labels = torch.randint(0, 64, (2, seq_len, seq_len))
     pde_labels = torch.randint(0, 64, (2, seq_len, seq_len))
     plddt_labels = torch.randint(0, 50, (2, seq_len))
@@ -427,7 +426,6 @@ def test_alphafold3():
         template_mask = template_mask,
         atom_pos = atom_pos,
         residue_atom_indices = residue_atom_indices,
-        distance_labels = distance_labels,
         pae_labels = pae_labels,
         pde_labels = pde_labels,
         plddt_labels = plddt_labels,