always atom resolution for trunk distogram

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -3392,7 +3392,6 @@ def __init__(
         dim_pairwise = 128,
         num_dist_bins = 38,
         dim_atom = 128,
-        atom_resolution = False
     ):
         super().__init__()
 
@@ -3404,10 +3403,7 @@ def __init__(
         # atom resolution
         # for now, just embed per atom distances, sum to atom features, project to pairwise dimension
 
-        self.atom_resolution = atom_resolution
-
-        if atom_resolution:
-            self.atom_feats_to_pairwise = LinearNoBiasThenOuterSum(dim_atom, dim_pairwise)
+        self.atom_feats_to_pairwise = LinearNoBiasThenOuterSum(dim_atom, dim_pairwise)
 
         # tensor typing
 
@@ -3417,16 +3413,12 @@ def __init__(
     def forward(
         self,
         pairwise_repr: Float['b n n d'],
-        molecule_atom_lens: Int['b n'] | None = None,
-        atom_feats: Float['b m {self.da}'] | None = None,
+        molecule_atom_lens: Int['b n'],
+        atom_feats: Float['b m {self.da}'],
     ) -> Float['b l n n'] | Float['b l m m']:
 
-        if self.atom_resolution:
-            assert exists(molecule_atom_lens)
-            assert exists(atom_feats)
-
-            pairwise_repr = batch_repeat_interleave_pairwise(pairwise_repr, molecule_atom_lens)
-            pairwise_repr = pairwise_repr + self.atom_feats_to_pairwise(atom_feats)
+        pairwise_repr = batch_repeat_interleave_pairwise(pairwise_repr, molecule_atom_lens)
+        pairwise_repr = pairwise_repr + self.atom_feats_to_pairwise(atom_feats)
 
         logits = self.to_distogram_logits(symmetrize(pairwise_repr))
 
@@ -4989,7 +4981,6 @@ def __init__(
         lddt_mask_other_cutoff = 15.,
         augment_kwargs: dict = dict(),
         stochastic_frame_average = False,
-        distogram_atom_resolution = False,
         checkpoint_input_embedding = False,
         checkpoint_trunk_pairformer = False,
         checkpoint_distogram_head = False,
@@ -5171,13 +5162,10 @@ def __init__(
 
         assert len(distance_bins_tensor) == num_dist_bins, '`distance_bins` must have a length equal to the `num_dist_bins` passed in'
 
-        self.distogram_atom_resolution = distogram_atom_resolution
-
         self.distogram_head = DistogramHead(
             dim_pairwise = dim_pairwise,
             dim_atom = dim_atom,
             num_dist_bins = num_dist_bins,
-            atom_resolution = distogram_atom_resolution,
         )
 
         # lddt related
@@ -5679,15 +5667,7 @@ def forward(
         if not exists(distance_labels) and atom_pos_given and exists(distogram_atom_indices):
 
             distogram_pos = atom_pos
-
-            if not self.distogram_atom_resolution:
-                # molecule_pos = einx.get_at('b [m] c, b n -> b n c', atom_pos, distogram_atom_indices)
-
-                distogram_atom_indices = repeat(distogram_atom_indices, 'b n -> b n c', c = distogram_pos.shape[-1])
-                molecule_pos = distogram_pos = distogram_pos.gather(1, distogram_atom_indices)
-                distogram_mask = valid_distogram_mask
-            else:
-                distogram_mask = atom_mask
+            distogram_mask = atom_mask
 
             distogram_dist = torch.cdist(distogram_pos, distogram_pos, p = 2)
             distance_labels = distance_to_bins(distogram_dist, self.distance_bins)
@@ -5700,9 +5680,7 @@ def forward(
         if exists(distance_labels):
 
             distogram_mask = pairwise_mask
-
-            if self.distogram_atom_resolution:
-                distogram_mask = to_pairwise_mask(atom_mask)
+            distogram_mask = to_pairwise_mask(atom_mask)
 
             distance_labels = torch.where(distogram_mask, distance_labels, ignore)
 

alphafold3_pytorch/mocks.py

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

pyproject.toml

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

tests/test_af3.py

@@ -511,27 +511,31 @@ def test_input_embedder():
     )
 
 def test_distogram_head():
+    molecule_atom_lens = torch.ones((2, 16)).long()
+    atom_feats = torch.randn(2, 16, 128)
     pairwise_repr = torch.randn(2, 16, 16, 128)
 
     distogram_head = DistogramHead(dim_pairwise = 128)
 
-    logits = distogram_head(pairwise_repr)
+    logits = distogram_head(
+        pairwise_repr,
+        atom_feats = atom_feats,
+        molecule_atom_lens = molecule_atom_lens
+    )
 
 @pytest.mark.parametrize('window_atompair_inputs', (True, False))
 @pytest.mark.parametrize('stochastic_frame_average', (True, False))
 @pytest.mark.parametrize('missing_atoms', (True, False))
 @pytest.mark.parametrize('calculate_pae', (True, False))
 @pytest.mark.parametrize('atom_transformer_intramolecular_attn', (True, False))
 @pytest.mark.parametrize('num_molecule_mods', (0, 4))
-@pytest.mark.parametrize('distogram_atom_resolution', (True, False))
 def test_alphafold3(
     window_atompair_inputs: bool,
     stochastic_frame_average: bool,
     missing_atoms: bool,
     calculate_pae: bool,
     atom_transformer_intramolecular_attn: bool,
     num_molecule_mods: int,
-    distogram_atom_resolution: bool
 ):
     seq_len = 16
     atom_seq_len = 32
@@ -622,7 +626,6 @@ def test_alphafold3(
             )
         ),
         stochastic_frame_average = stochastic_frame_average,
-        distogram_atom_resolution = distogram_atom_resolution
     )
 
     loss, breakdown = alphafold3(