allow for directed bonds, as defined by rdkit

Author: lucidrains

alphafold3_pytorch/inputs.py

@@ -180,6 +180,7 @@ class MoleculeInput:
     resolved_labels:            Int[' n'] | None = None
     add_atom_ids:               bool = False
     add_atompair_ids:           bool = False
+    directed_bonds:             bool = False
     extract_atom_feats_fn:      Callable[[Atom], Float['m dai']] = default_extract_atom_feats_fn
     extract_atompair_feats_fn:  Callable[[Mol], Float['m m dapi']] = default_extract_atompair_feats_fn
 
@@ -247,6 +248,8 @@ def molecule_to_atom_input(
 
     if i.add_atompair_ids:
         atom_bond_index = {symbol: (idx + 1) for idx, symbol in enumerate(ATOM_BONDS)}
+        num_atom_bond_types = len(atom_bond_index)
+
         other_index = len(ATOM_BONDS) + 1
 
         atompair_ids = torch.zeros(total_atoms, total_atoms).long()
@@ -284,7 +287,17 @@ def molecule_to_atom_input(
                 bond_type = bond.GetBondType()
                 bond_id = atom_bond_index.get(bond_type, other_index) + 1
 
-                updates.extend([bond_id, bond_id])
+                # default to symmetric bond type (undirected atom bonds)
+
+                bond_to = bond_from = bond_id
+
+                # if allowing for directed bonds, assume num_atompair_embeds = (2 * num_atom_bond_types) + 1
+                # offset other edge by num_atom_bond_types
+
+                if i.directed_bonds:
+                    bond_from += num_atom_bond_types
+
+                updates.extend([bond_to, bond_from])
 
             coordinates = tensor(coordinates).long()
             updates = tensor(updates).long()
@@ -386,6 +399,7 @@ class Alphafold3Input:
     add_atom_ids:               bool = False
     add_atompair_ids:           bool = False
     add_output_atompos_indices: bool = True
+    directed_bonds:             bool = False
     extract_atom_feats_fn:      Callable[[Atom], Float['m dai']] = default_extract_atom_feats_fn
     extract_atompair_feats_fn:  Callable[[Mol], Float['m m dapi']] = default_extract_atompair_feats_fn
 
@@ -833,6 +847,7 @@ def get_num_atoms_per_chain(chains: List[List[Mol]]) -> List[int]:
         atom_parent_ids = atom_parent_ids,
         add_atom_ids = i.add_atom_ids,
         add_atompair_ids = i.add_atompair_ids,
+        directed_bonds = i.directed_bonds,
         extract_atom_feats_fn = i.extract_atom_feats_fn,
         extract_atompair_feats_fn = i.extract_atompair_feats_fn
     )

pyproject.toml

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

tests/test_input.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 
 from alphafold3_pytorch import (
@@ -23,7 +24,8 @@ def test_tensor_reverse_complement():
     rc = reverse_complement_tensor(seq)
     assert torch.allclose(reverse_complement_tensor(rc), seq)
 
-def test_alphafold3_input():
+@pytest.mark.parametrize('directed_bonds', (False, True))
+def test_alphafold3_input(directed_bonds):
 
     alphafold3_input = Alphafold3Input(
         proteins = ['MLEICLKLVGCKSKKGLSSSSSCYLEEALQRPVASDF', 'MGKCRGLRTARKLRSHRRDQKWHDKQYKKAHLGTALKANPFGGASHAKGIVLEKVGVEAKQPNSAIRKCVRVQLIKNGKKITAFVPNDGCLNFIEENDEVLVAGFGRKGHAVGDIPGVRFKVVKVANVSLLALYKGKKERPRS'],
@@ -35,18 +37,21 @@ def test_alphafold3_input():
         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
+        add_atompair_ids = True,
+        directed_bonds = directed_bonds
     )
 
     batched_atom_input = alphafold3_inputs_to_batched_atom_input(alphafold3_input)
 
     # feed it into alphafold3
 
+    num_atom_bond_types = (6 * (2 if directed_bonds else 1))
+
     alphafold3 = Alphafold3(
         dim_atom_inputs = 3,
         dim_atompair_inputs = 1,
         num_atom_embeds = 47,
-        num_atompair_embeds = 6 + 1,
+        num_atompair_embeds = num_atom_bond_types + 1, # 0 is for no bond
         atoms_per_window = 27,
         dim_template_feats = 44,
         num_dist_bins = 38,