architecture is there

Author: lucidrains

README.md

@@ -1,6 +1,6 @@
 <img src="./alphafold3.png" width="500px"></img>
 
-## Alphafold 3 - Pytorch (wip)
+## Alphafold 3 - Pytorch
 
 Implementation of <a href="https://www.nature.com/articles/s41586-024-07487-w">Alphafold 3</a> in Pytorch
 

alphafold3_pytorch/alphafold3.py

@@ -3077,6 +3077,13 @@ def forward(
 
         assert exists(residue_atom_lens) or exists(atom_mask)
 
+        # if atompair inputs are not windowed, window it
+
+        is_atompair_inputs_windowed = atompair_inputs.ndim == 5
+
+        if not is_atompair_inputs_windowed:
+            atompair_inputs = full_pairwise_repr_to_windowed(atompair_inputs, window_size = self.atoms_per_window)
+
         # handle atom mask
 
         total_atoms = residue_atom_lens.sum(dim = -1)

pyproject.toml

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

tests/test_af3.py

@@ -26,7 +26,8 @@
 
 from alphafold3_pytorch.alphafold3 import (
     mean_pool_with_lens,
-    repeat_consecutive_with_lens
+    repeat_consecutive_with_lens,
+    full_pairwise_repr_to_windowed
 )
 
 def test_mean_pool_with_lens():
@@ -383,17 +384,25 @@ def test_distogram_head():
 
     logits = distogram_head(pairwise_repr)
 
-
-def test_alphafold3():
+@pytest.mark.parametrize('window_atompair_inputs', (True, False))
+def test_alphafold3(
+    window_atompair_inputs: bool
+):
     seq_len = 16
+    atoms_per_window = 27
+
     residue_atom_lens = torch.randint(1, 3, (2, seq_len))
     atom_seq_len = residue_atom_lens.sum(dim = -1).amax()
 
     token_bond = torch.randint(0, 2, (2, seq_len, seq_len)).bool()
 
     atom_inputs = torch.randn(2, atom_seq_len, 77)
+
     atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
 
+    if window_atompair_inputs:
+        atompair_inputs = full_pairwise_repr_to_windowed(atompair_inputs, window_size = atoms_per_window)
+
     additional_residue_feats = torch.randn(2, seq_len, 10)
 
     template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
@@ -412,6 +421,7 @@ def test_alphafold3():
 
     alphafold3 = Alphafold3(
         dim_atom_inputs = 77,
+        atoms_per_window = atoms_per_window,
         dim_template_feats = 44,
         num_dist_bins = 38,
         confidence_head_kwargs = dict(