first pass through molecule or atom level pae label derivation for training confidence head

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -5142,14 +5142,19 @@ def __init__(
             num_dist_bins = num_dist_bins
         )
 
-        # pae bins
+        # pae related bins and modules
 
         pae_bins_tensor = Tensor(pae_bins)
         self.register_buffer('pae_bins', pae_bins_tensor)
         num_pae_bins = len(pae_bins)
 
+        self.rigid_from_three_points = RigidFrom3Points()
+        self.compute_alignment_error = ComputeAlignmentError()
+
         # confidence head
 
+        self.confidence_head_atom_resolution = confidence_head_atom_resolution
+
         self.confidence_head = ConfidenceHead(
             dim_single_inputs = dim_single_inputs,
             atompair_dist_bins = distance_bins,
@@ -5286,7 +5291,6 @@ def forward(
         num_sample_steps: int | None = None,
         atom_pos: Float['b m 3'] | None = None,
         distance_labels: Int['b n n'] | Int['b m m'] | None = None,
-        pae_labels: Int['b n n'] | Int['b m m'] | None = None,
         pde_labels: Int['b n n'] | Int['b m m'] | None = None,
         plddt_labels: Int['b n'] | Int['b m'] | None = None,
         resolved_labels: Int['b n'] | Int['b m'] | None = None,
@@ -5544,7 +5548,7 @@ def forward(
 
         atom_pos_given = exists(atom_pos)
 
-        confidence_head_labels = (pae_labels, pde_labels, plddt_labels, resolved_labels)
+        confidence_head_labels = (atom_indices_for_frame, pde_labels, plddt_labels, resolved_labels)
         all_labels = (distance_labels, *confidence_head_labels)
 
         has_labels = any([*map(exists, all_labels)])
@@ -5592,14 +5596,17 @@ def forward(
                 mask = mask,
                 return_pae_logits = True
             )
-            if return_distogram_head_logits:
-                distogram_head_logits = self.distogram_head(pairwise.clone().detach())
-                return (
-                    sampled_atom_pos,
-                    confidence_head_logits,
-                    distogram_head_logits,
-                )
-            return sampled_atom_pos, confidence_head_logits
+
+            if not return_distogram_head_logits:
+                return sampled_atom_pos, confidence_head_logits
+
+            distogram_head_logits = self.distogram_head(pairwise.clone().detach())
+
+            return (
+                sampled_atom_pos,
+                confidence_head_logits,
+                distogram_head_logits,
+            )
 
         # if being forced to return loss, but do not have sufficient information to return losses, just return 0
 
@@ -5621,6 +5628,8 @@ def forward(
 
         # distogram head
 
+        molecule_pos = None
+
         if not exists(distance_labels) and atom_pos_given and exists(distogram_atom_indices):
             # molecule_pos = einx.get_at('b [m] c, b n -> b n c', atom_pos, distogram_atom_indices)
 
@@ -5668,8 +5677,10 @@ def forward(
                     additional_molecule_feats,
                     is_molecule_types,
                     molecule_atom_indices,
+                    molecule_pos,
+                    distogram_atom_indices,
+                    atom_indices_for_frame,
                     molecule_atom_lens,
-                    pae_labels,
                     pde_labels,
                     plddt_labels,
                     resolved_labels,
@@ -5692,8 +5703,10 @@ def forward(
                         additional_molecule_feats,
                         is_molecule_types,
                         molecule_atom_indices,
+                        molecule_pos,
+                        distogram_atom_indices,
+                        atom_indices_for_frame,
                         molecule_atom_lens,
-                        pae_labels,
                         pde_labels,
                         plddt_labels,
                         resolved_labels
@@ -5742,6 +5755,41 @@ def forward(
                 return_denoised_pos = True,
             )
 
+        # determine pae labels if possible
+
+        pae_labels = None
+        ch_atom_res = self.confidence_head_atom_resolution
+
+        if atom_pos_given and exists(atom_indices_for_frame):
+
+            denoised_molecule_pos = None
+
+            if not ch_atom_res:
+                assert exists(molecule_pos), '`distogram_atom_indices` must be passed in for calculating non-atomic PAE labels'
+                denoised_molecule_pos = denoised_atom_pos.gather(1, distogram_atom_indices)
+
+            three_atoms = einx.get_at('b [m] c, b n three -> three b n c', atom_pos, atom_indices_for_frame)
+            pred_three_atoms = einx.get_at('b [m] c, b n three -> three b n c', denoised_atom_pos, atom_indices_for_frame)
+
+            # compute frames
+
+            frames, _ = self.rigid_from_three_points(three_atoms)
+            pred_frames, _ = self.rigid_from_three_points(pred_three_atoms)
+
+            # align error
+
+            align_error = self.compute_alignment_error(
+                denoised_atom_pos if ch_atom_res else denoised_molecule_pos,
+                atom_pos if ch_atom_res else molecule_pos,
+                pred_frames,
+                frames,
+                molecule_atom_lens = molecule_atom_lens
+            )
+
+            # calculate pae labels as alignment error binned to 64 (0 - 32A)
+
+            pae_labels = distance_to_bins(align_error, self.pae_bins)
+
         # confidence head
 
         should_call_confidence_head = any([*map(exists, confidence_head_labels)])

alphafold3_pytorch/inputs.py

@@ -193,8 +193,8 @@ class AtomInput:
     missing_atom_mask:          Bool[' m'] | None = None
     molecule_atom_indices:      Int[' n'] | None = None
     distogram_atom_indices:     Int[' n'] | None = None
+    atom_indices_for_frame:     Int['n 3'] | None = None
     distance_labels:            Int['n n'] | None = None
-    pae_labels:                 Int['n n'] | None = None
     pde_labels:                 Int['n n'] | None = None
     plddt_labels:               Int[' n'] | None = None
     resolved_labels:            Int[' n'] | None = None
@@ -227,8 +227,8 @@ class BatchedAtomInput:
     missing_atom_mask:          Bool['b m'] | None = None
     molecule_atom_indices:      Int['b n'] | None = None
     distogram_atom_indices:     Int['b n'] | None = None
+    atom_indices_for_frame:     Int['b n 3'] | None = None
     distance_labels:            Int['b n n'] | None = None
-    pae_labels:                 Int['b n n'] | None = None
     pde_labels:                 Int['b n n'] | None = None
     plddt_labels:               Int['b n'] | None = None
     resolved_labels:            Int['b n'] | None = None
@@ -444,6 +444,7 @@ class MoleculeInput:
     is_molecule_mod:            Bool['n num_mods'] | Bool[' n'] | None = None
     molecule_atom_indices:      List[int | None] | None = None
     distogram_atom_indices:     List[int | None] | None = None
+    atom_indices_for_frame:     List[Tuple[int, int, int] | None] | None = None
     missing_atom_indices:       List[Int[' _'] | None] | None = None
     missing_token_indices:      List[Int[' _'] | None] | None = None
     atom_parent_ids:            Int[' m'] | None = None
@@ -454,7 +455,6 @@ class MoleculeInput:
     template_mask:              Bool[' t'] | None = None
     msa_mask:                   Bool[' s'] | None = None
     distance_labels:            Int['n n'] | None = None
-    pae_labels:                 Int['n n'] | None = None
     pde_labels:                 Int[' n'] | None = None
     resolved_labels:            Int[' n'] | None = None
     chains:                     Tuple[int | None, int | None] | None = (None, None)
@@ -715,6 +715,14 @@ def molecule_to_atom_input(mol_input: MoleculeInput) -> AtomInput:
     if is_molecule_mod.ndim == 1:
         is_molecule_mod = rearrange(is_molecule_mod, 'n -> n 1')
 
+    # handle `atom_indices_for_frame` for the PAE
+
+    atom_indices_for_frame = i.atom_indices_for_frame
+
+    if exists(atom_indices_for_frame):
+        atom_indices_for_frame = [default(indices, (-1, -1, -1)) for indices in i.atom_indices_for_frame]
+        atom_indices_for_frame = tensor(atom_indices_for_frame)
+
     # atom input
 
     atom_input = AtomInput(
@@ -724,6 +732,7 @@ def molecule_to_atom_input(mol_input: MoleculeInput) -> AtomInput:
         molecule_ids=i.molecule_ids,
         molecule_atom_indices=i.molecule_atom_indices,
         distogram_atom_indices=i.distogram_atom_indices,
+        atom_indices_for_frame=atom_indices_for_frame,
         is_molecule_mod=is_molecule_mod,
         msa=i.msa,
         templates=i.templates,
@@ -773,7 +782,6 @@ class MoleculeLengthMoleculeInput:
     template_mask:              Bool[' t'] | None = None
     msa_mask:                   Bool[' s'] | None = None
     distance_labels:            Int['n n'] | None = None
-    pae_labels:                 Int['n n'] | None = None
     pde_labels:                 Int[' n'] | None = None
     resolved_labels:            Int[' n'] | None = None
     chains:                     Tuple[int | None, int | None] | None = (None, None)
@@ -1194,7 +1202,6 @@ class Alphafold3Input:
     template_mask:              Bool[' t'] | None = None
     msa_mask:                   Bool[' s'] | None = None
     distance_labels:            Int['n n'] | None = None
-    pae_labels:                 Int['n n'] | None = None
     pde_labels:                 Int[' n'] | None = None
     resolved_labels:            Int[' n'] | None = None
     chains:                     Tuple[int | None, int | None] | None = (None, None)

alphafold3_pytorch/mocks.py

@@ -78,7 +78,6 @@ def __getitem__(self, idx):
         distogram_atom_indices = molecule_atom_lens - 1
 
         distance_labels = torch.randint(0, 37, (seq_len, seq_len))
-        pae_labels = torch.randint(0, 64, (seq_len, seq_len))
         pde_labels = torch.randint(0, 64, (seq_len, seq_len))
         plddt_labels = torch.randint(0, 50, (seq_len,))
         resolved_labels = torch.randint(0, 2, (seq_len,))
@@ -104,7 +103,6 @@ def __getitem__(self, idx):
             molecule_atom_indices = molecule_atom_indices,
             distogram_atom_indices = distogram_atom_indices,
             distance_labels = distance_labels,
-            pae_labels = pae_labels,
             pde_labels = pde_labels,
             plddt_labels = plddt_labels,
             resolved_labels = resolved_labels,

pyproject.toml

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

tests/test_af3.py

@@ -611,7 +611,6 @@ def test_alphafold3(
     molecule_atom_indices = molecule_atom_lens - 1
 
     label_len = atom_seq_len if confidence_head_atom_resolution else seq_len
-    pae_labels = torch.randint(0, 64, (2, label_len, label_len))
     pde_labels = torch.randint(0, 64, (2, label_len, label_len))
     plddt_labels = torch.randint(0, 50, (2, label_len))
     resolved_labels = torch.randint(0, 2, (2, label_len))
@@ -680,7 +679,6 @@ def test_alphafold3(
         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,
         resolved_labels = resolved_labels,
@@ -724,7 +722,6 @@ def test_alphafold3_without_msa_and_templates():
     distogram_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))
@@ -777,7 +774,6 @@ def test_alphafold3_without_msa_and_templates():
         atom_pos = atom_pos,
         distogram_atom_indices = distogram_atom_indices,
         distance_labels = distance_labels,
-        pae_labels = pae_labels,
         pde_labels = pde_labels,
         plddt_labels = plddt_labels,
         resolved_labels = resolved_labels,
@@ -803,7 +799,6 @@ def test_alphafold3_force_return_loss():
     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))
@@ -845,7 +840,6 @@ def test_alphafold3_force_return_loss():
         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,
@@ -888,7 +882,6 @@ def test_alphafold3_force_return_loss_with_confidence_logits():
     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))
@@ -930,7 +923,6 @@ def test_alphafold3_force_return_loss_with_confidence_logits():
         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,
@@ -996,7 +988,6 @@ def test_alphafold3_with_atom_and_bond_embeddings():
     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))
     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))
@@ -1022,7 +1013,6 @@ def test_alphafold3_with_atom_and_bond_embeddings():
         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