fix an issue with missing_atom_indices due to ligands being one token per atom

Author: lucidrains

alphafold3_pytorch/inputs.py

@@ -301,11 +301,11 @@ def molecule_to_atom_input(
 
         assert len(molecules) == len(i.missing_atom_indices), f'{len(i.missing_atom_indices)} missing atom indices does not match the number of molecules given ({len(molecules)})'
 
-        missing_atom_indices: List[Int[' _']] = [default(indices, torch.empty((0,), dtype = torch.long)) for indices in i.missing_atom_indices]
-
+        missing_atom_indices: List[Int[' _']] = []
         missing_atom_mask: List[Bool[' _']] = []
 
-        for num_atoms, mol_missing_atom_indices in zip(all_num_atoms, missing_atom_indices):
+        for num_atoms, mol_missing_atom_indices, is_ligand in zip(all_num_atoms, i.missing_atom_indices, i.is_molecule_types[:, -1]):
+            mol_missing_atom_indices = default(mol_missing_atom_indices, torch.empty((0,), dtype = torch.long))
 
             mol_miss_atom_mask = torch.zeros(num_atoms, dtype = torch.bool)
 
@@ -314,8 +314,14 @@ def molecule_to_atom_input(
 
             missing_atom_mask.append(mol_miss_atom_mask)
 
-        missing_atom_mask = torch.cat(missing_atom_mask)
+            if not is_ligand:
+                missing_atom_indices.append(mol_missing_atom_indices)
+            else:
+                for is_missing_atom_in_ligand in mol_miss_atom_mask:
+                    index = tensor([0] if is_missing_atom_in_ligand else [], dtype = torch.long)
+                    missing_atom_indices.append(index)
 
+        missing_atom_mask = torch.cat(missing_atom_mask)
         missing_atom_indices = pad_sequence(missing_atom_indices, batch_first = True, padding_value = -1)
 
     # handle maybe atompair embeds

pyproject.toml

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

tests/test_input.py

@@ -87,12 +87,14 @@ def test_atompos_input():
 
     mock_atompos = [
         torch.randn(5, 3),   # alanine has 5 non-hydrogen atoms
-        torch.randn(4, 3)    # glycine has 4 non-hydrogen atoms
+        torch.randn(4, 3),   # glycine has 4 non-hydrogen atoms
+        torch.randn(3, 3)    # ligand has 3 carbons
     ]
 
     train_alphafold3_input = Alphafold3Input(
         proteins = [contrived_protein],
-        missing_atom_indices = [[1, 2], None],
+        missing_atom_indices = [[1, 2], None, [0, 1]],
+        ligands = ['CCC'],
         atom_pos = mock_atompos
     )