Standardize all distogram binning logic (#213)

* Update README.md

* Update template_parsing.py

* Update alphafold3.py

* Update model_utils.py

* Update trainer_with_atom_dataset_created_from_pdb.yaml

* Update alphafold3.yaml

* Update trainer.yaml

* Update trainer_with_atom_dataset.yaml

* Update trainer_with_atom_dataset.yaml

* Update alphafold3.yaml

* Update trainer.yaml

* Update trainer_with_pdb_dataset.yaml

* Update trainer_with_pdb_dataset_and_weighted_sampling.yaml

* Update training_with_pdb_dataset.yaml

* Update test_dataloading.py

* Update test_input.py

* Update test_af3.py

* Update test_trainer.py

* Update test_af3.py

* Update mocks.py

* Update alphafold3.py

* Update alphafold3.py

Author: amorehead

README.md

@@ -202,7 +202,7 @@ alphafold3 = Alphafold3(
     dim_atompair_inputs = 5,
     atoms_per_window = 27,
     dim_template_feats = 108,
-    num_dist_bins = 38,
+    num_dist_bins = 64,
     num_molecule_mods = 0,
     confidence_head_kwargs = dict(
         pairformer_depth = 1

alphafold3_pytorch/alphafold3.py

@@ -4218,7 +4218,7 @@ def __init__(
         self,
         *,
         dim_pairwise = 128,
-        num_dist_bins = 38,
+        num_dist_bins = 64,
         dim_atom = 128,
         atom_resolution = False,
         checkpoint = False,
@@ -5236,11 +5236,7 @@ class ComputeModelSelectionScore(Module):
     def __init__(
         self,
         eps: float = 1e-8,
-        dist_breaks: Float[" dist_break"] = torch.linspace(  
-            2.3125,
-            21.6875,
-            37,
-        ),
+        dist_breaks: Float[" dist_break"] = torch.linspace(2, 22, 63),
         nucleic_acid_cutoff: float = 30.0,
         other_cutoff: float = 15.0,
         contact_mask_threshold: float = 8.0,
@@ -5891,7 +5887,7 @@ def __init__(
         num_atom_embeds: int | None = None,
         num_atompair_embeds: int | None = None,
         num_molecule_mods: int | None = DEFAULT_NUM_MOLECULE_MODS,
-        distance_bins: List[float] = torch.linspace(3, 20, 38).float().tolist(),
+        distance_bins: List[float] = torch.linspace(2, 22, 64).float().tolist(),  # NOTE: in paper, they reuse AF2's setup of having 64 bins from 2 to 22
         pae_bins: List[float] = torch.linspace(0.5, 32, 64).float().tolist(),
         pde_bins: List[float] = torch.linspace(0.5, 32, 64).float().tolist(),
         ignore_index = -1,

alphafold3_pytorch/data/template_parsing.py

@@ -121,16 +121,16 @@ def parse_m8(
             template_release_date = extract_mmcif_metadata_field(
                 template_mmcif_object, "release_date"
             )
-            template_biomol = _from_mmcif_object(
-                template_mmcif_object, chain_ids=set(template_chain)
-            )
             if not (
                 exists(template_cutoff_date)
                 and datetime.strptime(template_release_date, "%Y-%m-%d") <= template_cutoff_date
             ):
                 continue
             elif not exists(template_cutoff_date):
                 pass
+            template_biomol = _from_mmcif_object(
+                template_mmcif_object, chain_ids=set(template_chain)
+            )
             if len(template_biomol.atom_positions):
                 template_biomols.append((template_biomol, template_type))
         except Exception as e:
@@ -148,7 +148,7 @@ def _extract_template_features(
     query_chemtype: List[str],
     num_restype_classes: int = 32,
     num_distogram_bins: int = 39,
-    distance_bins: List[float] = torch.linspace(3.25, 50.75, 38).float(),
+    distance_bins: List[float] = torch.linspace(3.25, 50.75, 39).float(),
     verbose: bool = False,
 ) -> Dict[str, Any]:
     """Parse atom positions in the target structure and align with the query.
@@ -181,9 +181,9 @@ def _extract_template_features(
         f"Mapping length {len(mapping)} must match query sequence length {len(query_sequence)} "
         f"and query chemtype length {len(query_chemtype)}."
     )
-    assert num_distogram_bins == len(distance_bins) + 1, (
+    assert num_distogram_bins == len(distance_bins), (
         f"Number of distance bins {num_distogram_bins} must match the length of distance bins "
-        f"{len(distance_bins)} plus one."
+        f"{len(distance_bins)}."
     )
 
     all_atom_positions = template_biomol.atom_positions

alphafold3_pytorch/mocks.py

@@ -86,7 +86,7 @@ def __getitem__(self, idx):
         molecule_atom_indices += atom_offsets
         distogram_atom_indices += atom_offsets
 
-        distance_labels = torch.randint(0, 37, (seq_len, seq_len))
+        distance_labels = torch.randint(0, 64, (seq_len, seq_len))
         resolved_labels = torch.randint(0, 2, (atom_seq_len,))
 
         majority_asym_id = asym_id.mode().values.item()

alphafold3_pytorch/utils/model_utils.py

@@ -38,27 +38,31 @@ def default_lambda_lr_fn(steps: int) -> float:
 
 @typecheck
 def distance_to_dgram(
-    distance: Float['... dist'],
-    bins: Float[' bins'],
+    distance: Float["... dist"],  # type: ignore
+    bins: Float[" bins"],  # type: ignore
     return_labels: bool = False,
-) -> Int['... dist'] | Int['... dist bins']:
-    """
-    converting from distance to discrete bins, for distance_labels and pae_labels
-    using the same logic as openfold
+) -> Int["... dist"] | Int["... dist bins"]:  # type: ignore
+    """Converting from distance to discrete bins, e.g., for distance_labels and pae_labels using
+    the same logic as OpenFold.
+
+    :param distance: The distance tensor.
+    :param bins: The bins tensor.
+    :param return_labels: Whether to return the labels.
+    :return: The one-hot bins tensor or the bin labels.
     """
 
-    distance = distance ** 2
+    distance = distance.abs()
 
-    bins = F.pad(bins ** 2, (0, 1), value = float('inf'))
+    bins = F.pad(bins, (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)
+        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.argmax(dim=-1)
 
     return one_hot
 

tests/configs/trainer_with_atom_dataset_created_from_pdb.yaml

@@ -11,7 +11,7 @@ model:
     dim_template_model: 8
     atoms_per_window: 27
     dim_template_feats: 108
-    num_dist_bins: 38
+    num_dist_bins: 64
     ignore_index: -1
     num_dist_bins: null
     num_plddt_bins: 50

tests/configs/trainer_with_pdb_dataset.yaml

@@ -11,7 +11,7 @@ model:
     dim_template_model: 8
     atoms_per_window: 27
     dim_template_feats: 108
-    num_dist_bins: 38
+    num_dist_bins: 64
     ignore_index: -1
     num_dist_bins: null
     num_plddt_bins: 50

tests/configs/trainer_with_pdb_dataset_and_weighted_sampling.yaml

@@ -11,7 +11,7 @@ model:
     dim_template_model: 8
     atoms_per_window: 27
     dim_template_feats: 108
-    num_dist_bins: 38
+    num_dist_bins: 64
     ignore_index: -1
     num_dist_bins: null
     num_plddt_bins: 50

tests/configs/training_with_pdb_dataset.yaml

@@ -17,7 +17,7 @@ model:
     dim_template_model: 8
     atoms_per_window: 27
     dim_template_feats: 108
-    num_dist_bins: 38
+    num_dist_bins: 64
     ignore_index: -1
     num_dist_bins: null
     num_plddt_bins: 50

tests/test_af3.py

@@ -698,7 +698,7 @@ def test_alphafold3(
         dim_token = 8,
         atoms_per_window = atoms_per_window,
         dim_template_feats = 108,
-        num_dist_bins = 38,
+        num_dist_bins = 64,
         num_molecule_mods = num_molecule_mods,
         confidence_head_kwargs = dict(
             pairformer_depth = 1
@@ -804,7 +804,7 @@ def test_alphafold3_without_msa_and_templates():
     alphafold3 = Alphafold3(
         dim_atom_inputs = 77,
         dim_template_feats = 108,
-        num_dist_bins = 38,
+        num_dist_bins = 64,
         num_molecule_mods = 0,
         checkpoint_trunk_pairformer = True,
         checkpoint_diffusion_module = True,
@@ -871,7 +871,7 @@ def test_alphafold3_force_return_loss():
     alphafold3 = Alphafold3(
         dim_atom_inputs = 77,
         dim_template_feats = 108,
-        num_dist_bins = 38,
+        num_dist_bins = 64,
         num_molecule_mods = 0,
         confidence_head_kwargs = dict(
             pairformer_depth = 1
@@ -953,7 +953,7 @@ def test_alphafold3_force_return_loss_with_confidence_logits():
     alphafold3 = Alphafold3(
         dim_atom_inputs = 77,
         dim_template_feats = 108,
-        num_dist_bins = 38,
+        num_dist_bins = 64,
         num_molecule_mods = 0,
         confidence_head_kwargs = dict(
             pairformer_depth = 1
@@ -1170,7 +1170,7 @@ def test_model_selection_score():
     atom_mask = torch.randint(0, 2, (atom_pos_true.shape[:-1])).type_as(atom_pos_true).bool()
     tok_repr_atm_mask = torch.randint(0, 2, (batch_size, seq_len)).bool()
 
-    dist_logits = torch.randn(batch_size, 38, seq_len, seq_len)
+    dist_logits = torch.randn(batch_size, 64, seq_len, seq_len)
     pde_logits = torch.randn(batch_size, 64, seq_len, seq_len)
 
     chain_length = [random.randint(seq_len // 4, seq_len //2) 
@@ -1222,7 +1222,7 @@ def test_model_selection_score_end_to_end():
         dim_token = 8,
         atoms_per_window = 27,
         dim_template_feats = 108,
-        num_dist_bins = 38,
+        num_dist_bins = 64,
         confidence_head_kwargs = dict(
             pairformer_depth = 1
         ),
@@ -1427,7 +1427,7 @@ def test_readme2():
         dim_atompair_inputs = 5,
         atoms_per_window = 27,
         dim_template_feats = 108,
-        num_dist_bins = 38,
+        num_dist_bins = 64,
         num_molecule_mods = 0,
         confidence_head_kwargs = dict(
             pairformer_depth = 1