improvise a solution for atom resolution confidence heads, without atom attention

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -2853,10 +2853,10 @@ def forward(
 # confidence head
 
 class ConfidenceHeadLogits(NamedTuple):
-    pae: Float['b pae n n'] | None
-    pde: Float['b pde n n']
-    plddt: Float['b plddt n']
-    resolved: Float['b 2 n']
+    pae: Float['b pae n n'] | Float['b pae m m'] |  None
+    pde: Float['b pde n n'] | Float['b pde m m']
+    plddt: Float['b plddt n'] | Float['b plddt m']
+    resolved: Float['b 2 n'] | Float['b 2 m']
 
 class ConfidenceHead(Module):
     """ Algorithm 31 """
@@ -2866,6 +2866,8 @@ def __init__(
         self,
         *,
         dim_single_inputs,
+        atom_resolution = False,  # @amorehead discovers that the public api has per-atom resolution confidences. improvise a solution
+        dim_atom = 128,
         atompair_dist_bins: List[float],
         dim_single = 384,
         dim_pairwise = 128,
@@ -2918,6 +2920,19 @@ def __init__(
             Rearrange('b ... l -> b l ...')
         )
 
+        # 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_single = LinearNoBias(dim_atom, dim_single)
+            self.atom_feats_to_pairwise = LinearNoBiasThenOuterSum(dim_atom, dim_pairwise)
+
+        # tensor typing
+
+        self.da = dim_atom
+
     @typecheck
     def forward(
         self,
@@ -2927,7 +2942,9 @@ def forward(
         pairwise_repr: Float['b n n dp'],
         pred_atom_pos: Float['b n 3'] | Float['b m 3'],
         molecule_atom_indices: Int['b n'] | None = None,
+        molecule_atom_lens: Int['b n'] | None = None,
         mask: Bool['b n'] | None = None,
+        atom_feats: Float['b m {self.da}'] | None = None,
         return_pae_logits = True
 
     ) -> ConfidenceHeadLogits:
@@ -2938,16 +2955,24 @@ def forward(
 
         is_atom_seq = pred_atom_pos.shape[-2] > single_inputs_repr.shape[-2]
 
-        assert not is_atom_seq or exists(molecule_atom_indices)
+        # handle atom resolution vs not
+
+        if self.atom_resolution:
+            assert exists(atom_feats), 'atom_feats must be passed in if atom_resolution is turned on for ConfidenceHead'
+            assert is_atom_seq, '`pred_atom_pos` must be passed in with atomic length'
+            assert exists(molecule_atom_lens)
 
         if is_atom_seq:
-            pred_atom_pos = einx.get_at('b [m] c, b n -> b n c', pred_atom_pos, molecule_atom_indices)
+            assert exists(molecule_atom_indices), 'molecule_atom_indices must be passed into ConfidenceHead if pred_atom_pos is atomic length'
+            pred_molecule_pos = einx.get_at('b [m] c, b n -> b n c', pred_atom_pos, molecule_atom_indices)
+        else:
+            pred_molecule_pos = pred_atom_pos
 
         # interatomic distances - embed and add to pairwise
 
-        interatom_dist = torch.cdist(pred_atom_pos, pred_atom_pos, p = 2)
+        intermolecule_dist = torch.cdist(pred_molecule_pos, pred_molecule_pos, p = 2)
 
-        dist_from_dist_bins = einx.subtract('b m dist, dist_bins -> b m dist dist_bins', interatom_dist, self.atompair_dist_bins).abs()
+        dist_from_dist_bins = einx.subtract('b m dist, dist_bins -> b m dist dist_bins', intermolecule_dist, self.atompair_dist_bins).abs()
         dist_bin_indices = dist_from_dist_bins.argmin(dim = -1)
         pairwise_repr = pairwise_repr + self.dist_bin_pairwise_embed(dist_bin_indices)
 
@@ -2959,6 +2984,27 @@ def forward(
             mask = mask
         )
 
+        # handle maybe atom level resolution
+
+        if self.atom_resolution:
+            single_repr = repeat_consecutive_with_lens(single_repr, molecule_atom_lens)
+
+            pairwise_repr = repeat_consecutive_with_lens(pairwise_repr, molecule_atom_lens)
+
+            molecule_atom_lens = repeat(molecule_atom_lens, 'b ... -> (b r) ...', r = pairwise_repr.shape[1])
+            pairwise_repr, ps = pack_one(pairwise_repr, '* n d')
+            pairwise_repr = repeat_consecutive_with_lens(pairwise_repr, molecule_atom_lens)
+            pairwise_repr = unpack_one(pairwise_repr, ps, '* n d')
+
+            interatomic_dist = torch.cdist(pred_atom_pos, pred_atom_pos, p = 2)
+
+            dist_from_dist_bins = einx.subtract('b m dist, dist_bins -> b m dist dist_bins', interatomic_dist, self.atompair_dist_bins).abs()
+            dist_bin_indices = dist_from_dist_bins.argmin(dim = -1)
+            pairwise_repr = pairwise_repr + self.dist_bin_pairwise_embed(dist_bin_indices)
+
+            single_repr = single_repr + self.atom_feats_to_single(atom_feats)
+            pairwise_repr = pairwise_repr + self.atom_feats_to_pairwise(atom_feats)
+
         # to logits
 
         symmetric_pairwise_repr = pairwise_repr + rearrange(pairwise_repr, 'b i j d -> b j i d')

pyproject.toml

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

tests/test_af3.py

@@ -388,6 +388,41 @@ def test_confidence_head():
         mask = mask
     )
 
+def test_atom_resolution_confidence_head():
+    single_inputs_repr = torch.randn(2, 16, 77)
+    single_repr = torch.randn(2, 16, 384)
+    pairwise_repr = torch.randn(2, 16, 16, 128)
+    mask = torch.ones((2, 16)).bool()
+
+    atom_seq_len = 32
+    atom_feats = torch.randn(2, atom_seq_len, 64)
+    pred_atom_pos = torch.randn(2, atom_seq_len, 3)
+
+    molecule_atom_indices = torch.randint(0, atom_seq_len, (2, 16)).long()
+    molecule_atom_lens = torch.full((2, 16), 2).long()
+
+    confidence_head = ConfidenceHead(
+        dim_single_inputs = 77,
+        atom_resolution = True,
+        dim_atom = 64,
+        atompair_dist_bins = torch.linspace(3, 20, 37).tolist(),
+        dim_single = 384,
+        dim_pairwise = 128,
+    )
+
+    logits = confidence_head(
+        single_inputs_repr = single_inputs_repr,
+        single_repr = single_repr,
+        atom_feats = atom_feats,
+        molecule_atom_indices = molecule_atom_indices,
+        molecule_atom_lens = molecule_atom_lens,
+        pairwise_repr = pairwise_repr,
+        pred_atom_pos = pred_atom_pos,
+        mask = mask
+    )
+
+    assert logits.pde.shape[-1] == atom_seq_len
+
 def test_input_embedder():
 
     molecule_atom_lens = torch.randint(0, 3, (2, 16))