add is_molecule_mod to AtomInput and BatchAtomInput

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -54,6 +54,7 @@
     IS_LIGAND,
     IS_METAL_ION,
     NUM_MOLECULE_IDS,
+    DEFAULT_NUM_MOLECULE_MODS,
     ADDITIONAL_MOLECULE_FEATS
 )
 
@@ -4519,6 +4520,7 @@ def __init__(
         self.w = atoms_per_window
         self.dapi = self.dim_atompair_inputs
         self.dai = self.dim_atom_inputs
+        self.num_mods = num_molecule_mods
 
     @property
     def device(self):
@@ -4605,7 +4607,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 {self.w} {self.w*2}'] | None = None,
-        is_molecule_mod: Bool['b n num_mods'] | None = None,
+        is_molecule_mod: Bool['b n {self.num_mods}'] | None = None,
         atom_mask: Bool['b m'] | None = None,
         missing_atom_mask: Bool['b m'] | None = None,
         atom_parent_ids: Int['b m'] | None = None,

alphafold3_pytorch/inputs.py

@@ -79,6 +79,7 @@
 MOLECULE_METAL_ION_ID = MOLECULE_GAP_ID + 1
 NUM_MOLECULE_IDS = len(HUMAN_AMINO_ACIDS) + len(RNA_NUCLEOTIDES) + len(DNA_NUCLEOTIDES) + 2
 
+DEFAULT_NUM_MOLECULE_MODS = 5
 ADDITIONAL_MOLECULE_FEATS = 5
 
 CCD_COMPONENTS_FILEPATH = os.path.join('data', 'ccd_data', 'components.cif')
@@ -140,7 +141,8 @@ class AtomInput:
     atompair_inputs:            Float['m m dapi'] | Float['nw w (w*2) dapi']
     additional_molecule_feats:  Int[f'n {ADDITIONAL_MOLECULE_FEATS}']
     is_molecule_types:          Bool[f'n {IS_MOLECULE_TYPES}']
-    additional_token_feats:     Float[f'n dtf'] | None = None
+    is_molecule_mod:            Bool['n num_mods'] | None = None
+    additional_token_feats:     Float['n dtf'] | None = None
     templates:                  Float['t n n dt'] | None = None
     msa:                        Float['s n dm'] | None = None
     token_bonds:                Bool['n n'] | None = None
@@ -171,7 +173,8 @@ class BatchedAtomInput:
     atompair_inputs:            Float['b m m dapi'] | Float['b nw w (w*2) dapi']
     additional_molecule_feats:  Int[f'b n {ADDITIONAL_MOLECULE_FEATS}']
     is_molecule_types:          Bool[f'b n {IS_MOLECULE_TYPES}']
-    additional_token_feats:     Float[f'b n dtf'] | None = None
+    is_molecule_mod:            Bool['b n num_mods'] | None = None
+    additional_token_feats:     Float['b n dtf'] | None = None
     templates:                  Float['b t n n dt'] | None = None
     msa:                        Float['b s n dm'] | None = None
     token_bonds:                Bool['b n n'] | None = None
@@ -326,6 +329,7 @@ class MoleculeInput:
     is_molecule_types:          Bool[f'n {IS_MOLECULE_TYPES}']
     src_tgt_atom_indices:       Int['n 2']
     token_bonds:                Bool['n n']
+    is_molecule_mod:            Bool['n num_mods'] | None = None
     molecule_atom_indices:      List[int | None] | None = None
     distogram_atom_indices:     List[int | None] | None = None
     missing_atom_indices:       List[Int[' _'] | None] | None = None

alphafold3_pytorch/mocks.py

@@ -4,7 +4,11 @@
 import torch
 from torch.utils.data import Dataset
 from alphafold3_pytorch import AtomInput
-from alphafold3_pytorch.inputs import IS_MOLECULE_TYPES
+
+from alphafold3_pytorch.inputs import (
+    IS_MOLECULE_TYPES,
+    DEFAULT_NUM_MOLECULE_MODS
+)
 
 # mock dataset
 
@@ -13,11 +17,13 @@ def __init__(
         self,
         data_length,
         max_seq_len = 16,
-        atoms_per_window = 4
+        atoms_per_window = 4,
+        has_molecule_mods = False
     ):
         self.data_length = data_length
         self.max_seq_len = max_seq_len
         self.atoms_per_window = atoms_per_window
+        self.has_molecule_mods = has_molecule_mods
 
     def __len__(self):
         return self.data_length
@@ -33,6 +39,11 @@ def __getitem__(self, idx):
         additional_molecule_feats = torch.randint(0, 2, (seq_len, 5))
         additional_token_feats = torch.randn(seq_len, 2)
         is_molecule_types = torch.randint(0, 2, (seq_len, IS_MOLECULE_TYPES)).bool()
+
+        is_molecule_mod = None
+        if self.has_molecule_mods:
+            is_molecule_mod = torch.rand((seq_len, DEFAULT_NUM_MOLECULE_MODS)) < 0.05
+
         molecule_ids = torch.randint(0, 32, (seq_len,))
         token_bonds = torch.randint(0, 2, (seq_len, seq_len)).bool()
 
@@ -65,6 +76,7 @@ def __getitem__(self, idx):
             additional_molecule_feats = additional_molecule_feats,
             additional_token_feats = additional_token_feats,
             is_molecule_types = is_molecule_types,
+            is_molecule_mod = is_molecule_mod,
             templates = templates,
             template_mask = template_mask,
             msa = msa,

pyproject.toml

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