add token_bonds to AtomInput (should be renamed molecular bonds at some point, except for the fact they treat each atom of the ligand as a token..)

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -3221,7 +3221,7 @@ def forward(
         atom_ids: Int['b m'] | None = None,
         atompair_ids: Int['b m m'] | Int['b nw w1 w2'] | None = None,
         atom_mask: Bool['b m'] | None = None,
-        token_bond: Bool['b n n'] | None = None,
+        token_bonds: Bool['b n n'] | None = None,
         msa: Float['b s n d'] | None = None,
         msa_mask: Bool['b s'] | None = None,
         templates: Float['b t n n dt'] | None = None,
@@ -3322,21 +3322,21 @@ def forward(
 
         # token bond features
 
-        if exists(token_bond):
+        if exists(token_bonds):
             # well do some precautionary standardization
             # (1) mask out diagonal - token to itself does not count as a bond
             # (2) symmetrize, in case it is not already symmetrical (could also throw an error)
 
-            token_bond = token_bond | rearrange(token_bond, 'b i j -> b j i')
+            token_bonds = token_bonds | rearrange(token_bonds, 'b i j -> b j i')
             diagonal = torch.eye(seq_len, device = self.device, dtype = torch.bool)
-            token_bond = token_bond.masked_fill(diagonal, False)
+            token_bonds = token_bonds.masked_fill(diagonal, False)
         else:
             seq_arange = torch.arange(seq_len, device = self.device)
-            token_bond = einx.subtract('i, j -> i j', seq_arange, seq_arange).abs() == 1
+            token_bonds = einx.subtract('i, j -> i j', seq_arange, seq_arange).abs() == 1
 
-        token_bond_feats = self.token_bond_to_pairwise_feat(token_bond.float())
+        token_bonds_feats = self.token_bond_to_pairwise_feat(token_bonds.float())
 
-        pairwise_init = pairwise_init + token_bond_feats
+        pairwise_init = pairwise_init + token_bonds_feats
 
         # molecule mask and pairwise mask
 

alphafold3_pytorch/inputs.py

@@ -16,6 +16,7 @@ class AtomInput(TypedDict):
     additional_molecule_feats:  Float['n 9']
     templates:                  Float['t n n dt']
     msa:                        Float['s n dm']
+    token_bonds:                Bool['n n'] | None
     atom_ids:                   Int['m'] | None
     atompair_ids:               Int['m m'] | Int['nw w (w*2)'] | None
     template_mask:              Bool['t'] | None
@@ -36,6 +37,7 @@ class BatchedAtomInput(TypedDict):
     additional_molecule_feats:  Float['b n 9']
     templates:                  Float['b t n n dt']
     msa:                        Float['b s n dm']
+    token_bonds:                Bool['b n n'] | None
     atom_ids:                   Int['b m'] | None
     atompair_ids:               Int['b m m'] | Int['b nw w (w*2)'] | None
     template_mask:              Bool['b t'] | None

pyproject.toml

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

tests/test_af3.py

@@ -422,7 +422,7 @@ def test_alphafold3(
     molecule_atom_lens = torch.randint(1, 3, (2, seq_len))
     atom_seq_len = molecule_atom_lens.sum(dim = -1).amax()
 
-    token_bond = torch.randint(0, 2, (2, seq_len, seq_len)).bool()
+    token_bonds = torch.randint(0, 2, (2, seq_len, seq_len)).bool()
 
     atom_inputs = torch.randn(2, atom_seq_len, 77)
 
@@ -480,7 +480,7 @@ def test_alphafold3(
         molecule_atom_lens = molecule_atom_lens,
         atompair_inputs = atompair_inputs,
         additional_molecule_feats = additional_molecule_feats,
-        token_bond = token_bond,
+        token_bonds = token_bonds,
         msa = msa,
         msa_mask = msa_mask,
         templates = template_feats,

tests/test_trainer.py

@@ -48,6 +48,7 @@ def __getitem__(self, idx):
         molecule_atom_lens = torch.randint(1, self.atoms_per_window, (seq_len,))
         additional_molecule_feats = torch.randn(seq_len, 9)
         molecule_ids = torch.randint(0, 32, (seq_len,))
+        token_bonds = torch.randint(0, 2, (seq_len, seq_len)).bool()
 
         templates = torch.randn(2, seq_len, seq_len, 44)
         template_mask = torch.ones((2,)).bool()
@@ -73,6 +74,7 @@ def __getitem__(self, idx):
             atom_inputs = atom_inputs,
             atompair_inputs = atompair_inputs,
             molecule_ids = molecule_ids,
+            token_bonds = token_bonds,
             molecule_atom_lens = molecule_atom_lens,
             additional_molecule_feats = additional_molecule_feats,
             templates = templates,