comment and quick assert

Author: lucidrains

alphafold3_pytorch/inputs.py

@@ -740,6 +740,8 @@ def get_num_atoms_per_chain(chains: List[List[Mol]]) -> List[int]:
         if isinstance(atom_pos, list):
             atom_pos = torch.cat(atom_pos, dim = -2)
 
+        assert atom_pos.shape[-2] == total_atoms
+
         # to automatically reorder the atom positions back to canonical
 
         if i.add_output_atompos_indices:

pyproject.toml

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

tests/test_input.py

@@ -79,8 +79,8 @@ def test_atompos_input():
     contrived_protein = 'AG'
 
     mock_atompos = [
-        torch.randn(6, 3),
-        torch.randn(5, 3)
+        torch.randn(6, 3),   # alanine has 6 non-hydrogen atoms
+        torch.randn(5, 3)    # glycine has 5 non-hydrogen atoms
     ]
 
     train_alphafold3_input = Alphafold3Input(