add a hard validation for the indices under an env flag (#174)

add a hard validation for all indices under an env flag or `hard_validate` on Alphafold3.forward

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -59,6 +59,7 @@
     DEFAULT_NUM_MOLECULE_MODS,
     ADDITIONAL_MOLECULE_FEATS,
     BatchedAtomInput,
+    hard_validate_atom_indices_ascending
 )
 
 from alphafold3_pytorch.common.biomolecule import (
@@ -5347,7 +5348,8 @@ def forward(
         rollout_show_tqdm_pbar: bool = False,
         detach_when_recycling: bool = None,
         min_conf_resolution: float = 0.1,
-        max_conf_resolution: float = 4.0
+        max_conf_resolution: float = 4.0,
+        hard_validate: bool = False
     ) -> (
         Float['b m 3'] |
         Tuple[Float['b m 3'], ConfidenceHeadLogits] |
@@ -5363,6 +5365,13 @@ def forward(
         assert atom_inputs.shape[-1] == self.dim_atom_inputs, f'expected {self.dim_atom_inputs} for atom_inputs feature dimension, but received {atom_inputs.shape[-1]}'
         assert atompair_inputs.shape[-1] == self.dim_atompair_inputs, f'expected {self.dim_atompair_inputs} for atompair_inputs feature dimension, but received {atompair_inputs.shape[-1]}'
 
+        # hard validate when debug env variable is turned on
+
+        if hard_validate or IS_DEBUGGING:
+            maybe(hard_validate_atom_indices_ascending)(distogram_atom_indices, 'distogram_atom_indices')
+            maybe(hard_validate_atom_indices_ascending)(molecule_atom_indices, 'molecule_atom_indices')
+            maybe(hard_validate_atom_indices_ascending)(atom_indices_for_frame, 'atom_indices_for_frame')
+
         # soft validate
 
         valid_molecule_atom_mask = valid_atom_len_mask = molecule_atom_lens >= 0
@@ -5385,8 +5394,6 @@ def forward(
 
         assert exists(molecule_atom_lens) or exists(atom_mask)
 
-        # hard validate when debug env variable is turned on
-
         if IS_DEBUGGING:
             assert (molecule_atom_lens >= 0).all(), 'molecule_atom_lens must be greater or equal to 0'
 

alphafold3_pytorch/inputs.py

@@ -166,6 +166,33 @@ def inner(x, *args, **kwargs):
         return fn(x, *args, **kwargs)
     return inner
 
+# validation functions
+
+def hard_validate_atom_indices_ascending(
+    indices: Int['b n'] | Int['b n 3'],
+    error_msg_field: str = 'indices'
+):
+    # will do a hard validate
+    # asserting if any of the indices that are not -1 (missing) are identical or descending
+    # this will cover 'distogram_atom_indices', 'molecule_atom_indices', and 'atom_indices_for_frame'
+
+    if indices.ndim == 2:
+        indices = rearrange(indices, '... -> ... 1')
+
+    for batch_index, sample_indices in enumerate(indices):
+
+        all_present = (sample_indices >= 0).all(dim = -1)
+        present_indices = sample_indices[all_present]
+
+        # relaxed assumption that if all -1 or only one molecule, it passes the test
+
+        if present_indices.numel() <= 1:
+            continue
+
+        difference = einx.subtract('n i, n j -> n (i j)', present_indices[1:], present_indices[:-1])
+
+        assert (difference >= 0).all(), f'detected invalid {error_msg_field} for in a batch: {present_indices}'
+
 # atom level, what Alphafold3 accepts
 
 UNCOLLATABLE_ATOM_INPUT_FIELDS = {'filepath'}

alphafold3_pytorch/mocks.py

@@ -9,6 +9,8 @@
     DEFAULT_NUM_MOLECULE_MODS
 )
 
+from alphafold3_pytorch.utils.model_utils import exclusive_cumsum
+
 # mock dataset
 
 class MockAtomDataset(Dataset):
@@ -37,6 +39,8 @@ def __getitem__(self, idx):
         atompair_inputs = torch.randn(atom_seq_len, atom_seq_len, 5)
 
         molecule_atom_lens = torch.randint(1, self.atoms_per_window, (seq_len,))
+        atom_offsets = exclusive_cumsum(molecule_atom_lens)
+
         additional_molecule_feats = torch.randint(0, 2, (seq_len, 5))
         additional_token_feats = torch.randn(seq_len, 2)
         is_molecule_types = torch.randint(0, 2, (seq_len, IS_MOLECULE_TYPES)).bool()
@@ -77,6 +81,9 @@ def __getitem__(self, idx):
         molecule_atom_indices = molecule_atom_lens - 1
         distogram_atom_indices = molecule_atom_lens - 1
 
+        molecule_atom_indices += atom_offsets
+        distogram_atom_indices += atom_offsets
+
         distance_labels = torch.randint(0, 37, (atom_seq_len, atom_seq_len))
         resolved_labels = torch.randint(0, 2, (atom_seq_len,))
 

pyproject.toml

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

tests/test_af3.py

@@ -1,5 +1,6 @@
 import os
 os.environ['TYPECHECK'] = 'True'
+os.environ['DEBUG'] = 'True'
 
 import pytest
 import random
@@ -63,6 +64,8 @@
     default_extract_atompair_feats_fn
 )
 
+from alphafold3_pytorch.utils.model_utils import exclusive_cumsum
+
 DATA_TEST_PDB_ID = '721p'
 
 def test_atom_ref_pos_to_atompair_inputs():
@@ -538,11 +541,13 @@ def test_alphafold3(
     num_molecule_mods: int,
 ):
     seq_len = 16
-    atom_seq_len = 32
     atoms_per_window = 27
 
     molecule_atom_indices = torch.randint(0, 2, (2, seq_len)).long()
-    molecule_atom_lens = torch.full((2, seq_len), 2).long()
+    molecule_atom_lens = torch.full((2, seq_len), 3).long()
+
+    atom_seq_len = molecule_atom_lens.sum(dim = -1).amax()
+    atom_offset = exclusive_cumsum(molecule_atom_lens)
 
     token_bonds = torch.randint(0, 2, (2, seq_len, seq_len)).bool()
 
@@ -564,7 +569,8 @@ def test_alphafold3(
 
     atom_indices_for_frame = None
     if calculate_pae:
-        atom_indices_for_frame = repeat(torch.arange(3), 'c -> b n c', b = 2, n = seq_len)
+        atom_indices_for_frame = repeat(torch.arange(3), 'c -> b n c', b = 2, n = seq_len).clone()
+        atom_indices_for_frame += atom_offset[..., None]
 
     missing_atom_mask = None
     if missing_atoms:
@@ -586,6 +592,13 @@ def test_alphafold3(
 
     resolved_labels = torch.randint(0, 2, (2, atom_seq_len))
 
+    # offset indices correctly
+
+    distogram_atom_indices += atom_offset
+    molecule_atom_indices += atom_offset
+
+    # alphafold3
+
     alphafold3 = Alphafold3(
         dim_atom_inputs = 77,
         dim_pairwise = 8,
@@ -750,8 +763,11 @@ def test_alphafold3_force_return_loss():
     seq_len = 16
     atom_seq_len = 32
 
-    molecule_atom_indices = torch.randint(0, 2, (2, seq_len)).long()
     molecule_atom_lens = torch.full((2, seq_len), 2).long()
+    atom_offsets = exclusive_cumsum(molecule_atom_lens)
+
+    molecule_atom_indices = torch.randint(0, 2, (2, seq_len)).long()
+    molecule_atom_indices += atom_offsets
 
     atom_inputs = torch.randn(2, atom_seq_len, 77)
     atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
@@ -761,7 +777,9 @@ def test_alphafold3_force_return_loss():
     molecule_ids = torch.randint(0, 32, (2, seq_len))
 
     atom_pos = torch.randn(2, atom_seq_len, 3)
+
     distogram_atom_indices = molecule_atom_lens - 1
+    distogram_atom_indices += atom_offsets
 
     resolved_labels = torch.randint(0, 2, (2, atom_seq_len))
 
@@ -828,8 +846,11 @@ def test_alphafold3_force_return_loss_with_confidence_logits():
     seq_len = 16
     atom_seq_len = 32
 
-    molecule_atom_indices = torch.randint(0, 2, (2, seq_len)).long()
     molecule_atom_lens = torch.full((2, seq_len), 2).long()
+    atom_offsets = exclusive_cumsum(molecule_atom_lens)
+
+    molecule_atom_indices = torch.randint(0, 2, (2, seq_len)).long()
+    molecule_atom_indices += atom_offsets
 
     atom_inputs = torch.randn(2, atom_seq_len, 77)
     atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
@@ -840,6 +861,7 @@ def test_alphafold3_force_return_loss_with_confidence_logits():
 
     atom_pos = torch.randn(2, atom_seq_len, 3)
     distogram_atom_indices = molecule_atom_lens - 1
+    distogram_atom_indices += atom_offsets
 
     resolved_labels = torch.randint(0, 2, (2, atom_seq_len))
 
@@ -916,8 +938,11 @@ def test_alphafold3_with_atom_and_bond_embeddings():
     seq_len = 16
     atom_seq_len = 32
 
-    molecule_atom_indices = torch.randint(0, 2, (2, seq_len)).long()
     molecule_atom_lens = torch.full((2, seq_len), 2).long()
+    atom_offset = exclusive_cumsum(molecule_atom_lens)
+
+    molecule_atom_indices = torch.randint(0, 2, (2, seq_len)).long()
+    molecule_atom_indices += atom_offset
 
     atom_ids = torch.randint(0, 7, (2, atom_seq_len))
     atompair_ids = torch.randint(0, 3, (2, atom_seq_len, atom_seq_len))
@@ -940,6 +965,7 @@ def test_alphafold3_with_atom_and_bond_embeddings():
 
     atom_pos = torch.randn(2, atom_seq_len, 3)
     distogram_atom_indices = molecule_atom_lens - 1 # last atom, as an example
+    distogram_atom_indices += atom_offset
 
     resolved_labels = torch.randint(0, 2, (2, atom_seq_len))