add ability to specify custom atom list

Author: lucidrains

alphafold3_pytorch/inputs.py

@@ -761,7 +761,7 @@ class MoleculeInput:
     directed_bonds: bool = False
     extract_atom_feats_fn: Callable[[Atom], Float["m dai"]] = default_extract_atom_feats_fn  # type: ignore
     extract_atompair_feats_fn: Callable[[Mol], Float["m m dapi"]] = default_extract_atompair_feats_fn  # type: ignore
-
+    custom_atoms: List[str]| None = None
 
 @typecheck
 def molecule_to_atom_input(mol_input: MoleculeInput) -> AtomInput:
@@ -803,7 +803,9 @@ def molecule_to_atom_input(mol_input: MoleculeInput) -> AtomInput:
     atom_ids = None
 
     if i.add_atom_ids:
-        atom_index = {symbol: i for i, symbol in enumerate(ATOMS)}
+        atom_list = default(i.custom_atoms, ATOMS)
+
+        atom_index = {symbol: i for i, symbol in enumerate(atom_list)}
 
         atom_ids = []
 
@@ -1101,6 +1103,7 @@ class MoleculeLengthMoleculeInput:
     directed_bonds: bool = False
     extract_atom_feats_fn: Callable[[Atom], Float["m dai"]] = default_extract_atom_feats_fn  # type: ignore
     extract_atompair_feats_fn: Callable[[Mol], Float["m m dapi"]] = default_extract_atompair_feats_fn  # type: ignore
+    custom_atoms: List[str]| None = None
 
 
 @typecheck
@@ -1264,7 +1267,9 @@ def molecule_lengthed_molecule_input_to_atom_input(
     atom_ids = None
 
     if i.add_atom_ids:
-        atom_index = {symbol: i for i, symbol in enumerate(ATOMS)}
+        atom_list = default(i.custom_atoms, ATOMS)
+
+        atom_index = {symbol: i for i, symbol in enumerate(atom_list)}
 
         atom_ids = []
 
@@ -1547,7 +1552,7 @@ class Alphafold3Input:
     directed_bonds: bool = False
     extract_atom_feats_fn: Callable[[Atom], Float["m dai"]] = default_extract_atom_feats_fn  # type: ignore
     extract_atompair_feats_fn: Callable[[Mol], Float["m m dapi"]] = default_extract_atompair_feats_fn  # type: ignore
-
+    custom_atoms: List[str] | None = None
 
 @typecheck
 def map_int_or_string_indices_to_mol(
@@ -1994,6 +1999,7 @@ def get_num_atoms_per_chain(chains: List[List[Mol]]) -> List[int]:
         directed_bonds=i.directed_bonds,
         extract_atom_feats_fn=i.extract_atom_feats_fn,
         extract_atompair_feats_fn=i.extract_atompair_feats_fn,
+        custom_atoms=i.custom_atoms
     )
 
     return molecule_input
@@ -3975,6 +3981,7 @@ def pdb_input_to_molecule_input(
         directed_bonds=i.directed_bonds,
         extract_atom_feats_fn=i.extract_atom_feats_fn,
         extract_atompair_feats_fn=i.extract_atompair_feats_fn,
+        custom_atoms=i.custom_atoms
     )
 
     return molecule_input

tests/test_input.py

@@ -20,6 +20,7 @@
 
 from alphafold3_pytorch.data.data_pipeline import *
 from alphafold3_pytorch.data.data_pipeline import make_mmcif_features
+
 from alphafold3_pytorch.common.biomolecule import (
     Biomolecule,
     _from_mmcif_object,
@@ -33,6 +34,7 @@
 from alphafold3_pytorch.data import mmcif_writing, mmcif_parsing
 
 from alphafold3_pytorch.life import (
+    ATOMS,
     reverse_complement,
     reverse_complement_tensor
 )
@@ -84,6 +86,8 @@ def test_alphafold3_input(
     directed_bonds
 ):
 
+    CUSTOM_ATOMS = list({*ATOMS, 'Na', 'Fe', 'Si', 'F', 'K'})
+
     alphafold3_input = Alphafold3Input(
         proteins = ['MLEICLKLVGCKSKKGLSSSSSCYLEEALQRPVASDF', 'MGKCRGLRTARKLRSHRRDQKWHDKQYKKAHLGTALKANPFGGASHAKGIVLEKVGVEAKQPNSAIRKCVRVQLIKNGKKITAFVPNDGCLNFIEENDEVLVAGFGRKGHAVGDIPGVRFKVVKVANVSLLALYKGKKERPRS'],
         ds_dna = ['ACGTT'],
@@ -95,7 +99,8 @@ def test_alphafold3_input(
         ligands = ['CC1=C(C=C(C=C1)NC(=O)C2=CC=C(C=C2)CN3CCN(CC3)C)NC4=NC=CC(=N4)C5=CN=CC=C5'],
         add_atom_ids = True,
         add_atompair_ids = True,
-        directed_bonds = directed_bonds
+        directed_bonds = directed_bonds,
+        custom_atoms = CUSTOM_ATOMS
     )
 
     batched_atom_input = alphafold3_inputs_to_batched_atom_input(alphafold3_input)
@@ -107,7 +112,7 @@ def test_alphafold3_input(
     alphafold3 = Alphafold3(
         dim_atom_inputs = 3,
         dim_atompair_inputs = 5,
-        num_atom_embeds = 47,
+        num_atom_embeds = len(CUSTOM_ATOMS),
         num_atompair_embeds = num_atom_bond_types + 1, # 0 is for no bond
         atoms_per_window = 27,
         dim_template_feats = 108,
@@ -187,17 +192,17 @@ def test_alphafold3_input_to_mmcif(tmp_path):
     """Test the Inference I/O Pipeline. This codifies the data_pipeline.py file used for training."""
 
     alphafold3_input = Alphafold3Input(
-    proteins = ['MLEICLKLVGCKSKKGLSSSSSCYLEEALQRPVASDF', 'MGKCRGLRTARKLRSHRRDQKWHDKQYKKAHLGTALKANPFGGASHAKGIVLEKVGVEAKQPNSAIRKCVRVQLIKNGKKITAFVPNDGCLNFIEENDEVLVAGFGRKGHAVGDIPGVRFKVVKVANVSLLALYKGKKERPRS'],
-    ds_dna = ['ACGTT'],
-    ds_rna = ['GCCAU', 'CCAGU'],
-    ss_dna = ['GCCTA'],
-    ss_rna = ['CGCAUA'],
-    metal_ions = ['Na', 'Na', 'Fe'],
-    misc_molecule_ids = ['Phospholipid'],
-    ligands = ['CC1=C(C=C(C=C1)NC(=O)C2=CC=C(C=C2)CN3CCN(CC3)C)NC4=NC=CC(=N4)C5=CN=CC=C5'],
-    add_atom_ids = True,
-    add_atompair_ids = True,
-    directed_bonds = True
+        proteins = ['MLEICLKLVGCKSKKGLSSSSSCYLEEALQRPVASDF', 'MGKCRGLRTARKLRSHRRDQKWHDKQYKKAHLGTALKANPFGGASHAKGIVLEKVGVEAKQPNSAIRKCVRVQLIKNGKKITAFVPNDGCLNFIEENDEVLVAGFGRKGHAVGDIPGVRFKVVKVANVSLLALYKGKKERPRS'],
+        ds_dna = ['ACGTT'],
+        ds_rna = ['GCCAU', 'CCAGU'],
+        ss_dna = ['GCCTA'],
+        ss_rna = ['CGCAUA'],
+        metal_ions = ['Na', 'Na', 'Fe'],
+        misc_molecule_ids = ['Phospholipid'],
+        ligands = ['CC1=C(C=C(C=C1)NC(=O)C2=CC=C(C=C2)CN3CCN(CC3)C)NC4=NC=CC(=N4)C5=CN=CC=C5'],
+        add_atom_ids = True,
+        add_atompair_ids = True,
+        directed_bonds = True
     )
 
     test_biomol = alphafold3_input_to_biomolecule(alphafold3_input, atom_positions=torch.randn(261, 47, 3).numpy())
@@ -328,7 +333,7 @@ def test_atompos_input():
             atom_encoder_depth = 1,
             token_transformer_depth = 1,
             atom_decoder_depth = 1,
-        )
+        ),
     )
 
     loss = alphafold3(**batched_atom_input.model_forward_dict())