add distogram loss off pairwise trunk for starters

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -2011,6 +2011,12 @@ def __init__(
         self.loss_confidence_weight = loss_confidence_weight
         self.loss_diffusion_weight = loss_diffusion_weight
 
+        self.register_buffer('dummy', torch.tensor(0), persistent = False)
+
+    @property
+    def device(self):
+        return self.dummy.device
+
     @typecheck
     def forward(
         self,
@@ -2048,6 +2054,7 @@ def forward(
         w = self.atoms_per_window
 
         mask = reduce(atom_mask, 'b (n w) -> b n', w = w, reduction = 'any')
+        pairwise_mask = einx.logical_and('b i, b j -> b i j', mask, mask)
 
         # init recycled single and pairwise
 
@@ -2107,6 +2114,22 @@ def forward(
         # otherwise will sample the atomic coordinates
 
         labels = (distance_labels, pae_labels, pde_labels, plddt_labels, resolved_labels)
-        return_loss = any([*filter(exists, labels)])
+        return_loss = any([*map(exists, labels)])
+
+        if not return_loss:
+            return torch.randn((*atom_inputs.shape[:2], 3), device = self.device)
+
+        # calculate all logits and losses
+
+        ignore = self.ignore_index
+
+        if exists(distance_labels):
+            distance_labels = torch.where(pairwise_mask, distance_labels, ignore)
+            distogram_logits = self.distogram_head(pairwise)
+            distogram_loss = F.cross_entropy(distogram_logits, distance_labels, ignore_index = ignore)
+
+        loss = (
+            distogram_loss * self.loss_distogram_weight
+        )
 
-        return torch.tensor(0.)
+        return loss

tests/test_readme.py

@@ -246,10 +246,13 @@ def test_alphafold3():
 
     msa = torch.randn(2, 7, seq_len, 64)
 
+    distance_labels = torch.randint(0, 38, (2, seq_len, seq_len))
+
     alphafold3 = Alphafold3(
         dim_atom_inputs = 77,
         dim_additional_residue_feats = 33,
-        dim_template_feats = 44
+        dim_template_feats = 44,
+        num_dist_bins = 38
     )
 
     loss = alphafold3(
@@ -260,7 +263,8 @@ def test_alphafold3():
         additional_residue_feats = additional_residue_feats,
         msa = msa,
         templates = template_feats,
-        template_mask = template_mask
+        template_mask = template_mask,
+        distance_labels = distance_labels
     )
 
-    print(loss)
+    loss.backward()