handle any type of molecule modifications (phospho, glycosy, methyl) all in one swoop, as a sparse boolean tensor

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -2967,6 +2967,7 @@ def __init__(
         num_molecule_types: int = 32,       # restype in additional residue information, apparently 32 (must be human amino acids + nucleotides + something else)
         num_atom_embeds: int | None = None,
         num_atompair_embeds: int | None = None,
+        num_molecule_mods: int | None = None,
         distance_bins: List[float] = torch.linspace(3, 20, 38).float().tolist(),
         ignore_index = -1,
         num_dist_bins: int | None = None,
@@ -3064,6 +3065,16 @@ def __init__(
         self.has_atom_embeds = has_atom_embeds
         self.has_atompair_embeds = has_atompair_embeds
 
+        # residue or nucleotide modifications
+
+        has_molecule_mod_embeds = num_molecule_mods > 0
+        self.num_molecule_mods = num_molecule_mods
+
+        if has_molecule_mod_embeds:
+            self.molecule_mod_embeds = nn.Embedding(num_molecule_mods, dim_single)
+
+        self.has_molecule_mod_embeds = has_molecule_mod_embeds
+
         # atoms per window
 
         self.atoms_per_window = atoms_per_window
@@ -3302,6 +3313,7 @@ def forward(
         additional_token_feats: Float['b n {self.dim_additional_token_feats}'] | None = None,
         atom_ids: Int['b m'] | None = None,
         atompair_ids: Int['b m m'] | Int['b nw w1 w2'] | None = None,
+        is_molecule_mod: Bool['b n {self.num_molecule_mods}'] | None = None,
         atom_mask: Bool['b m'] | None = None,
         atom_parent_ids: Int['b m'] | None = None,
         token_bonds: Bool['b n n'] | None = None,
@@ -3401,6 +3413,26 @@ def forward(
 
             atompair_feats = atompair_feats + atompair_embeds
 
+        # handle maybe molecule modifications
+
+        assert not (exists(is_molecule_mod) ^ self.has_molecule_mod_embeds), 'you either set `num_molecule_mods` and did not pass in `is_molecule_mod` or vice versa'
+
+        if self.has_molecule_mod_embeds:
+            single_init, seq_packed_shape = pack_one(single_init, '* ds')
+
+            is_molecule_mod, _ = pack_one(is_molecule_mod, '* mods')
+
+            if not is_molecule_mod.is_sparse:
+                is_molecule_mod = is_molecule_mod.to_sparse()
+
+            seq_indices, mod_id = is_molecule_mod.indices()
+            scatter_values = self.molecule_mod_embeds(mod_id)
+
+            seq_indices = repeat(seq_indices, 'n -> n ds', ds = single_init.shape[-1])
+            single_init = single_init.scatter_add(0, seq_indices, scatter_values)
+
+            single_init = unpack_one(single_init, seq_packed_shape, '* ds')
+
         # relative positional encoding
 
         relative_position_encoding = self.relative_position_encoding(

pyproject.toml

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

tests/test_af3.py

@@ -418,10 +418,12 @@ def test_distogram_head():
 @pytest.mark.parametrize('window_atompair_inputs', (True, False))
 @pytest.mark.parametrize('stochastic_frame_average', (True, False))
 @pytest.mark.parametrize('atom_transformer_intramolecular_attn', (True, False))
+@pytest.mark.parametrize('num_molecule_mods', (0, 5))
 def test_alphafold3(
     window_atompair_inputs: bool,
     stochastic_frame_average: bool,
-    atom_transformer_intramolecular_attn: bool
+    atom_transformer_intramolecular_attn: bool,
+    num_molecule_mods: int
 ):
     seq_len = 16
     atoms_per_window = 27
@@ -443,6 +445,10 @@ def test_alphafold3(
     is_molecule_types = torch.randint(0, 2, (2, seq_len, 4)).bool()
     molecule_ids = torch.randint(0, 32, (2, seq_len))
 
+    is_molecule_mod = None
+    if num_molecule_mods > 0:
+        is_molecule_mod = torch.zeros(2, seq_len, num_molecule_mods).uniform_(0, 1) < 0.1
+
     atom_parent_ids = None
 
     if atom_transformer_intramolecular_attn:
@@ -495,6 +501,7 @@ def test_alphafold3(
         atom_parent_ids = atom_parent_ids,
         atompair_inputs = atompair_inputs,
         is_molecule_types = is_molecule_types,
+        is_molecule_mod = is_molecule_mod,
         additional_molecule_feats = additional_molecule_feats,
         additional_token_feats = additional_token_feats,
         token_bonds = token_bonds,
@@ -519,6 +526,7 @@ def test_alphafold3(
         atom_inputs = atom_inputs,
         molecule_ids = molecule_ids,
         molecule_atom_lens = molecule_atom_lens,
+        is_molecule_mod = is_molecule_mod,
         atompair_inputs = atompair_inputs,
         is_molecule_types = is_molecule_types,
         additional_molecule_feats = additional_molecule_feats,