address #122

Author: lucidrains

README.md

@@ -174,7 +174,7 @@ batched_atom_input = alphafold3_inputs_to_batched_atom_input(train_alphafold3_in
 
 alphafold3 = Alphafold3(
     dim_atom_inputs = 3,
-    dim_atompair_inputs = 1,
+    dim_atompair_inputs = 5,
     atoms_per_window = 27,
     dim_template_feats = 44,
     num_dist_bins = 38,

alphafold3_pytorch/alphafold3.py

@@ -230,57 +230,6 @@ def pad_and_window(
     t = rearrange(t, 'b (n w) ... -> b n w ...', w = window_size)
     return t
 
-# to atompair input functions
-
-@typecheck
-def atom_ref_pos_to_atompair_inputs(
-    atom_ref_pos: Float['... m 3'],
-    atom_ref_space_uid: Int['... m'],
-) -> Float['... m m 5']:
-
-    # Algorithm 5 - lines 2-6
-    # allow for either batched or single
-
-    atom_ref_pos, unpack_one = pack_one(atom_ref_pos, '* m c')
-    atom_ref_space_uid, _ = pack_one(atom_ref_space_uid, '* m')
-
-    assert atom_ref_pos.shape[0] == atom_ref_space_uid.shape[0]
-
-    # line 2
-
-    pairwise_rel_pos = einx.subtract('b i c, b j c -> b i j c', atom_ref_pos, atom_ref_pos)
-
-    # line 3
-
-    same_ref_space_mask = einx.equal('b i, b j -> b i j', atom_ref_space_uid, atom_ref_space_uid)
-
-    # line 5 - pairwise inverse squared distance
-
-    atom_inv_square_dist = (1 + pairwise_rel_pos.norm(dim = -1, p = 2) ** 2) ** -1
-
-    # concat all into atompair_inputs for projection into atompair_feats within Alphafold3
-
-    atompair_inputs, _ = pack((
-        pairwise_rel_pos,
-        atom_inv_square_dist,
-        same_ref_space_mask.float(),
-    ), 'b i j *')
-
-    # mask out
-
-    atompair_inputs = einx.where(
-        'b i j, b i j dapi, -> b i j dapi',
-        same_ref_space_mask, atompair_inputs, 0.
-    )
-
-    # reconstitute optional batch dimension
-
-    atompair_inputs = unpack_one(atompair_inputs, '* i j dapi')
-
-    # return
-
-    return atompair_inputs
-
 # packed atom representation functions
 
 @typecheck

alphafold3_pytorch/inputs.py

@@ -12,6 +12,7 @@
 from typing import Any, Callable, List, Literal, Set, Tuple, Type
 
 import einx
+from einops import pack
 
 import numpy as np
 from numpy.lib.format import open_memmap
@@ -54,7 +55,6 @@
     reverse_complement_tensor,
 )
 
-
 from alphafold3_pytorch.tensor_typing import Bool, Float, Int, typecheck
 from alphafold3_pytorch.utils.data_utils import RESIDUE_MOLECULE_TYPE, get_residue_molecule_type
 from alphafold3_pytorch.utils.model_utils import exclusive_cumsum
@@ -298,6 +298,51 @@ def __getitem__(self, idx: int) -> AtomInput:
         file = self.files[idx]
         return file_to_atom_input(file)
 
+# atom reference position to atompair inputs
+# will be used in the `default_extract_atompair_feats_fn` below in MoleculeInput
+
+@typecheck
+def atom_ref_pos_to_atompair_inputs(
+    atom_ref_pos: Float['m 3'],
+    atom_ref_space_uid: Int['m'] | None = None,
+) -> Float['m m 5']:
+
+    # Algorithm 5 - lines 2-6
+
+    # line 2
+
+    pairwise_rel_pos = einx.subtract('i c, j c -> i j c', atom_ref_pos, atom_ref_pos)
+
+    # line 5 - pairwise inverse squared distance
+
+    atom_inv_square_dist = (1 + pairwise_rel_pos.norm(dim = -1, p = 2) ** 2) ** -1
+    
+    # line 3
+
+    if exists(atom_ref_space_uid):
+        same_ref_space_mask = einx.equal('i, j -> i j', atom_ref_space_uid, atom_ref_space_uid)
+    else:
+        same_ref_space_mask = torch.ones_like(atom_inv_square_dist).bool()
+
+    # concat all into atompair_inputs for projection into atompair_feats within Alphafold3
+
+    atompair_inputs, _ = pack((
+        pairwise_rel_pos,
+        atom_inv_square_dist,
+        same_ref_space_mask.float(),
+    ), 'i j *')
+
+    # mask out
+
+    atompair_inputs = einx.where(
+        'i j, i j dapi, -> i j dapi',
+        same_ref_space_mask, atompair_inputs, 0.
+    )
+
+    # return
+
+    return atompair_inputs
+
 # molecule input - accepting list of molecules as rdchem.Mol + the atomic lengths for how to pool into tokens
 
 def default_extract_atom_feats_fn(atom: Atom):
@@ -316,8 +361,7 @@ def default_extract_atompair_feats_fn(mol: Mol):
 
     all_atom_pos_tensor = tensor(all_atom_pos)
 
-    dist_matrix = torch.cdist(all_atom_pos_tensor, all_atom_pos_tensor)
-    return torch.stack((dist_matrix,), dim = -1)
+    return atom_ref_pos_to_atompair_inputs(all_atom_pos_tensor) # what they did in the paper, but can be overwritten
 
 @typecheck
 @dataclass

pyproject.toml

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

tests/configs/trainer_with_atom_dataset_created_from_pdb.yaml

@@ -7,7 +7,7 @@ model:
     dim_pairwise: 4
     dim_token: 4
     dim_atom_inputs: 3
-    dim_atompair_inputs: 1
+    dim_atompair_inputs: 5
     dim_template_model: 8
     atoms_per_window: 27
     dim_template_feats: 44

tests/configs/trainer_with_pdb_dataset.yaml

@@ -7,7 +7,7 @@ model:
     dim_pairwise: 4
     dim_token: 4
     dim_atom_inputs: 3
-    dim_atompair_inputs: 1
+    dim_atompair_inputs: 5
     dim_template_model: 8
     atoms_per_window: 27
     dim_template_feats: 44

tests/configs/trainer_with_pdb_dataset_and_weighted_sampling.yaml

@@ -7,7 +7,7 @@ model:
     dim_pairwise: 4
     dim_token: 4
     dim_atom_inputs: 3
-    dim_atompair_inputs: 1
+    dim_atompair_inputs: 5
     dim_template_model: 8
     atoms_per_window: 27
     dim_template_feats: 44

tests/configs/training_with_pdb_dataset.yaml

@@ -13,7 +13,7 @@ model:
     dim_pairwise: 4
     dim_token: 4
     dim_atom_inputs: 3
-    dim_atompair_inputs: 1
+    dim_atompair_inputs: 5
     dim_template_model: 8
     atoms_per_window: 27
     dim_template_feats: 44

tests/test_af3.py

@@ -40,12 +40,12 @@
 from alphafold3_pytorch.alphafold3 import (
     mean_pool_with_lens,
     repeat_consecutive_with_lens,
-    full_pairwise_repr_to_windowed,
-    atom_ref_pos_to_atompair_inputs
+    full_pairwise_repr_to_windowed
 )
 
 from alphafold3_pytorch.inputs import (
-    IS_MOLECULE_TYPES
+    IS_MOLECULE_TYPES,
+    atom_ref_pos_to_atompair_inputs
 )
 
 def test_atom_ref_pos_to_atompair_inputs():

tests/test_dataloading.py

@@ -41,7 +41,7 @@ def test_data_input():
 
     alphafold3 = Alphafold3(
         dim_atom_inputs=3,
-        dim_atompair_inputs=1,
+        dim_atompair_inputs=5,
         atoms_per_window=27,
         dim_template_feats=44,
         num_dist_bins=38,