#201 Move dgram calculation to model_utils, parametrize label returning (#208)

Author: vandrw

alphafold3_pytorch/alphafold3.py

@@ -75,6 +75,8 @@
     package_available,
 )
 
+from alphafold3_pytorch.utils.model_utils import distance_to_dgram
+
 from frame_averaging_pytorch import FrameAverage
 
 from taylor_series_linear_attention import TaylorSeriesLinearAttn
@@ -455,28 +457,6 @@ def batch_repeat_interleave_pairwise(
     pairwise = batch_repeat_interleave(pairwise, molecule_atom_lens)
     return unpack_one(pairwise)
 
-@typecheck
-def distance_to_bins(
-    distance: Float['... dist'],
-    bins: Float[' bins']
-) -> Int['... dist']:
-    """
-    converting from distance to discrete bins, for distance_labels and pae_labels
-    using the same logic as openfold
-    """
-
-    distance = distance ** 2
-
-    bins = F.pad(bins ** 2, (0, 1), value = float('inf'))
-    low, high = bins[:-1], bins[1:]
-
-    one_hot = (
-        einx.greater_equal('..., bin_low -> ... bin_low', distance, low) &
-        einx.less('..., bin_high -> ... bin_high', distance, high)
-    ).long()
-
-    return one_hot.argmax(dim = -1)
-
 # linear and outer sum
 # for single repr -> pairwise pattern throughout this architecture
 
@@ -4501,7 +4481,9 @@ def forward(
 
         intermolecule_dist = torch.cdist(pred_molecule_pos, pred_molecule_pos, p=2)
 
-        dist_bin_indices = distance_to_bins(intermolecule_dist, self.atompair_dist_bins)
+        dist_bin_indices = distance_to_dgram(
+            intermolecule_dist, self.atompair_dist_bins, return_labels=True
+        )
         pairwise_repr = pairwise_repr + self.dist_bin_pairwise_embed(dist_bin_indices)
 
         # pairformer stack
@@ -6734,7 +6716,9 @@ def forward(
                 distogram_mask = atom_mask
 
             distogram_dist = torch.cdist(distogram_pos, distogram_pos, p=2)
-            distance_labels = distance_to_bins(distogram_dist, self.distance_bins)
+            distance_labels = distance_to_dgram(
+                distogram_dist, self.distance_bins, return_labels = True
+            )
 
             # account for representative distogram atom missing from residue (-1 set on distogram_atom_indices field)
 
@@ -7014,9 +6998,9 @@ def forward(
                     mask=align_error_mask,
                 )
 
-                # calculate pae labels as alignment error binned to 64 (0 - 32A) (TODO: double-check correctness of `distance_to_bins`'s bin assignments)
+                # calculate pae labels as alignment error binned to 64 (0 - 32A)
 
-                pae_labels = distance_to_bins(align_error, self.pae_bins)
+                pae_labels = distance_to_dgram(align_error, self.pae_bins, return_labels = True)
 
                 # set ignore index for invalid molecules or frames
 
@@ -7058,7 +7042,7 @@ def forward(
                 # calculate pde labels as distance error binned to 64 (0 - 32A)
 
                 pde_dist = torch.abs(pde_pred_dist - pde_gt_dist)
-                pde_labels = distance_to_bins(pde_dist, self.pde_bins)
+                pde_labels = distance_to_dgram(pde_dist, self.pde_bins, return_labels = True)
 
                 # account for representative molecule atom missing from residue (-1 set on molecule_atom_indices field)
 

alphafold3_pytorch/data/template_parsing.py

@@ -24,7 +24,7 @@
 from alphafold3_pytorch.utils.model_utils import (
     ExpressCoordinatesInFrame,
     RigidFrom3Points,
-    distance_to_bins,
+    distance_to_dgram,
     get_frames_from_atom_pos,
 )
 from alphafold3_pytorch.tensor_typing import typecheck
@@ -367,17 +367,7 @@ def _extract_template_features(
     template_distogram_dist = torch.cdist(
         template_distogram_atom_positions, template_distogram_atom_positions, p=2
     )
-    template_distogram_dist_binned = distance_to_bins(template_distogram_dist, distance_bins)
-
-    template_distogram_dist_binned[
-        # NOTE: This assigns the last bin to distances greater than the maximum bin (e.g., > 50.75 Å).
-        template_distogram_dist
-        > distance_bins.max()
-    ] = (num_distogram_bins - 1)
-
-    template_distogram = F.one_hot(
-        template_distogram_dist_binned, num_classes=num_distogram_bins
-    ).float()
+    template_distogram = distance_to_dgram(template_distogram_dist, distance_bins)
 
     # Construct unit vectors.
     template_unit_vector = torch.zeros(

alphafold3_pytorch/utils/model_utils.py

@@ -37,20 +37,30 @@ def default_lambda_lr_fn(steps: int) -> float:
 
 
 @typecheck
-def distance_to_bins(
-    distance: Float["... dist"],  # type: ignore
-    bins: Float[" bins"],  # type: ignore
-) -> Int["... dist"]:  # type: ignore
-    """Convert from distance to discrete bins, e.g., for `distance_labels`.
-
-    :param distance: The distance tensor.
-    :param bins: The bins tensor.
-    :return: The discrete bins.
-    """
-    dist_from_dist_bins = einx.subtract(
-        "... dist, dist_bins -> ... dist dist_bins", distance, bins
-    ).abs()
-    return dist_from_dist_bins.argmin(dim=-1)
+def distance_to_dgram(
+    distance: Float['... dist'],
+    bins: Float[' bins'],
+    return_labels: bool = False,
+) -> Int['... dist']:
+    """
+    converting from distance to discrete bins, for distance_labels and pae_labels
+    using the same logic as openfold
+    """
+
+    distance = distance ** 2
+
+    bins = F.pad(bins ** 2, (0, 1), value = float('inf'))
+    low, high = bins[:-1], bins[1:]
+
+    one_hot = (
+        einx.greater_equal('..., bin_low -> ... bin_low', distance, low) &
+        einx.less('..., bin_high -> ... bin_high', distance, high)
+    ).long()
+
+    if return_labels:
+        return one_hot.argmax(dim = -1)
+
+    return one_hot
 
 
 @typecheck