make sure one can return sampled atom positions along with confidence head logits

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -3386,9 +3386,16 @@ def forward(
         return_loss_breakdown = False,
         return_loss: bool = None,
         return_present_sampled_atoms: bool = False,
+        return_confidence_head_logits: bool = False,
         num_rollout_steps: int = 20,
         rollout_show_tqdm_pbar: bool = False
-    ) -> Float['b m 3'] | Float['l 3'] | Float[''] | Tuple[Float[''], LossBreakdown]:
+    ) -> (
+        Float['b m 3'] |
+        Float['l 3'] |
+        Tuple[Float['b m 3'] | Float['l 3'], ConfidenceHeadLogits] |
+        Float[''] |
+        Tuple[Float[''], LossBreakdown]
+    ):
 
         atom_seq_len = atom_inputs.shape[-2]
 
@@ -3635,7 +3642,25 @@ def forward(
             if exists(missing_atom_mask) and return_present_sampled_atoms:
                 sampled_atom_pos = sampled_atom_pos[~missing_atom_mask]
 
-            return sampled_atom_pos
+            if not return_confidence_head_logits:
+                return sampled_atom_pos
+
+            # todo - handle missing atoms
+
+            assert exists(molecule_atom_indices)
+
+            pred_atom_pos = einx.get_at('b [m] c, b n -> b n c', sampled_atom_pos, molecule_atom_indices)
+
+            logits = self.confidence_head(
+                single_repr = single.detach(),
+                single_inputs_repr = single_inputs.detach(),
+                pairwise_repr = pairwise.detach(),
+                pred_atom_pos = pred_atom_pos.detach(),
+                mask = mask,
+                return_pae_logits = True
+            )
+
+            return sampled_atom_pos, logits
 
         # if being forced to return loss, but do not have sufficient information to return losses, just return 0
 

pyproject.toml

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

tests/test_af3.py

@@ -472,6 +472,7 @@ def test_alphafold3(
 
     atom_pos = torch.randn(2, atom_seq_len, 3)
     distogram_atom_indices = molecule_atom_lens - 1
+    molecule_atom_indices = molecule_atom_lens - 1
 
     pae_labels = torch.randint(0, 64, (2, seq_len, seq_len))
     pde_labels = torch.randint(0, 64, (2, seq_len, seq_len))
@@ -524,6 +525,7 @@ def test_alphafold3(
         template_mask = template_mask,
         atom_pos = atom_pos,
         distogram_atom_indices = distogram_atom_indices,
+        molecule_atom_indices = molecule_atom_indices,
         pae_labels = pae_labels,
         pde_labels = pde_labels,
         plddt_labels = plddt_labels,
@@ -630,6 +632,7 @@ def test_alphafold3_force_return_loss():
 
     atom_pos = torch.randn(2, atom_seq_len, 3)
     distogram_atom_indices = molecule_atom_lens - 1
+    molecule_atom_indices = molecule_atom_lens - 1
 
     distance_labels = torch.randint(0, 38, (2, seq_len, seq_len))
     pae_labels = torch.randint(0, 64, (2, seq_len, seq_len))
@@ -671,6 +674,7 @@ def test_alphafold3_force_return_loss():
         additional_token_feats = additional_token_feats,
         atom_pos = atom_pos,
         distogram_atom_indices = distogram_atom_indices,
+        molecule_atom_indices = molecule_atom_indices,
         distance_labels = distance_labels,
         pae_labels = pae_labels,
         pde_labels = pde_labels,
@@ -682,6 +686,91 @@ def test_alphafold3_force_return_loss():
 
     assert sampled_atom_pos.ndim == 3
 
+    loss, _ = alphafold3(
+        num_recycling_steps = 2,
+        atom_inputs = atom_inputs,
+        molecule_ids = molecule_ids,
+        molecule_atom_lens = molecule_atom_lens,
+        atompair_inputs = atompair_inputs,
+        is_molecule_types = is_molecule_types,
+        additional_molecule_feats = additional_molecule_feats,
+        additional_token_feats = additional_token_feats,
+        molecule_atom_indices = molecule_atom_indices,
+        return_loss_breakdown = True,
+        return_loss = True # force returning loss even if no labels given
+    )
+
+    assert loss == 0.
+
+def test_alphafold3_force_return_loss_with_confidence_logits():
+    seq_len = 16
+    molecule_atom_lens = torch.randint(1, 3, (2, seq_len))
+    atom_seq_len = molecule_atom_lens.sum(dim = -1).amax()
+
+    atom_inputs = torch.randn(2, atom_seq_len, 77)
+    atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
+    additional_molecule_feats = torch.randint(0, 2, (2, seq_len, 5))
+    additional_token_feats = torch.randn(2, seq_len, 2)
+    is_molecule_types = torch.randint(0, 2, (2, seq_len, IS_MOLECULE_TYPES)).bool()
+    molecule_ids = torch.randint(0, 32, (2, seq_len))
+
+    atom_pos = torch.randn(2, atom_seq_len, 3)
+    distogram_atom_indices = molecule_atom_lens - 1
+    molecule_atom_indices = molecule_atom_lens - 1
+
+    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))
+    resolved_labels = torch.randint(0, 2, (2, seq_len))
+
+    alphafold3 = Alphafold3(
+        dim_atom_inputs = 77,
+        dim_template_feats = 44,
+        num_dist_bins = 38,
+        confidence_head_kwargs = dict(
+            pairformer_depth = 1
+        ),
+        template_embedder_kwargs = dict(
+            pairformer_stack_depth = 1
+        ),
+        msa_module_kwargs = dict(
+            depth = 1
+        ),
+        pairformer_stack = dict(
+            depth = 2
+        ),
+        diffusion_module_kwargs = dict(
+            atom_encoder_depth = 1,
+            token_transformer_depth = 1,
+            atom_decoder_depth = 1,
+        ),
+    )
+
+    sampled_atom_pos, confidence_head_logits = alphafold3(
+        num_recycling_steps = 2,
+        atom_inputs = atom_inputs,
+        molecule_ids = molecule_ids,
+        molecule_atom_lens = molecule_atom_lens,
+        atompair_inputs = atompair_inputs,
+        is_molecule_types = is_molecule_types,
+        additional_molecule_feats = additional_molecule_feats,
+        additional_token_feats = additional_token_feats,
+        atom_pos = atom_pos,
+        distogram_atom_indices = distogram_atom_indices,
+        molecule_atom_indices = molecule_atom_indices,
+        distance_labels = distance_labels,
+        pae_labels = pae_labels,
+        pde_labels = pde_labels,
+        plddt_labels = plddt_labels,
+        resolved_labels = resolved_labels,
+        return_loss_breakdown = True,
+        return_loss = False, # force sampling even if labels are given
+        return_confidence_head_logits = True
+    )
+
+    assert sampled_atom_pos.ndim == 3
+
     loss, _ = alphafold3(
         num_recycling_steps = 2,
         atom_inputs = atom_inputs,
@@ -733,6 +822,7 @@ def test_alphafold3_with_atom_and_bond_embeddings():
 
     atom_pos = torch.randn(2, atom_seq_len, 3)
     distogram_atom_indices = molecule_atom_lens - 1 # last atom, as an example
+    molecule_atom_indices = molecule_atom_lens - 1
 
     distance_labels = torch.randint(0, 37, (2, seq_len, seq_len))
     pae_labels = torch.randint(0, 64, (2, seq_len, seq_len))
@@ -759,6 +849,7 @@ def test_alphafold3_with_atom_and_bond_embeddings():
         template_mask = template_mask,
         atom_pos = atom_pos,
         distogram_atom_indices = distogram_atom_indices,
+        molecule_atom_indices = molecule_atom_indices,
         distance_labels = distance_labels,
         pae_labels = pae_labels,
         pde_labels = pde_labels,