add atom pooling with fixed window size for eventual atom14 support

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -292,6 +292,30 @@ def mean_pool_with_lens(
     avg = einx.where('b n, b n d, -> b n d', mask, avg, 0.)
     return avg
 
+@typecheck
+def mean_pool_fixed_windows_with_mask(
+    feats: Float['b m d'],
+    mask: Bool['b m'],
+    window_size: int,
+    return_pooled_mask: bool = False
+) -> Float['b n d'] | Tuple[Float['b n'], Bool['b n']]:
+
+    seq_len = feats.shape[-2]
+    assert divisible_by(seq_len, window_size)
+
+    feats = einx.where('b m, b m d, -> b m d', mask, feats, 0.)
+
+    num = reduce(feats, 'b (n w) d -> b n d', 'sum', w = window_size)
+    den = reduce(mask.float(), 'b (n w) -> b n 1', 'sum', w = window_size)
+
+    avg = num / den.clamp(min = 1.)
+
+    if not return_pooled_mask:
+        return avg
+
+    pooled_mask = reduce(mask, 'b (n w) -> b n', 'any', w = window_size)
+    return avg, pooled_mask
+
 @typecheck
 def batch_repeat_interleave(
     feats: Float['b n ...'] | Bool['b n ...'] | Bool['b n'] | Int['b n'],

pyproject.toml

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

tests/test_af3.py

@@ -43,6 +43,7 @@
 
 from alphafold3_pytorch.alphafold3 import (
     mean_pool_with_lens,
+    mean_pool_fixed_windows_with_mask,
     batch_repeat_interleave,
     full_pairwise_repr_to_windowed,
     get_cid_molecule_type,
@@ -75,6 +76,14 @@ def test_mean_pool_with_lens():
 
     assert torch.allclose(pooled, torch.tensor([[[1.], [2.], [1.]]]))
 
+def test_mean_pool_with_mask():
+    seq = torch.tensor([[[1.], [100.], [1.], [2.], [2.], [100.], [1.], [1.], [100.]]])
+    mask = torch.tensor([[True, False, True, True, True, False, True, True, False]])
+
+    pooled = mean_pool_fixed_windows_with_mask(seq, mask, window_size = 3)
+
+    assert torch.allclose(pooled, torch.tensor([[[1.], [2.], [1.]]]))
+
 def test_batch_repeat_interleave():
     seq = torch.tensor([[[1.], [2.], [4.]], [[1.], [2.], [4.]]])
     lens = torch.tensor([[3, 4, 2], [2, 5, 1]]).long()