Add templates and atom-wise ligand frames for confidence head outputs (#195)

* Update README.md

* Update __init__.py

* Update data_pipeline.py

* Create kalign.py

* Update life.py

* Update mocks.py

* Create template_parsing.py

* Update alphafold3.py

* Update data_utils.py

* Update inputs.py

* Update model_utils.py

* Update test_af3.py

* Update test_dataloading.py

* Update test_input.py

* Create test_template_loading.py

* Update inputs.py

* Update template_parsing.py

* Update test_af3.py

* Update test_input.py

* Update test_trainer.py

* Update alphafold3.yaml

* Update trainer.yaml

* Update trainer_with_atom_dataset.yaml

* Update trainer_with_atom_dataset_created_from_pdb.yaml

* Update trainer_with_pdb_dataset.yaml

* Update trainer_with_pdb_dataset_and_weighted_sampling.yaml

* Update training.yaml

* Update training_with_pdb_dataset.yaml

* Update test_input.py

* Update inputs.py

Author: amorehead

README.md

@@ -42,7 +42,7 @@ A visualization of the molecules of life used in the repository can be seen and
 
 - <a href="https://github.com/xluo233">@xluo233</a> for contributing the confidence measures, clash penalty ranking, and sample ranking logic!
 
-- <a href="https://github.com/sj900">sj900</a> for integrating and testing the `WeightedPDBSampler` within the `PDBDataset` and for adding initial support for MSA parsing!
+- <a href="https://github.com/sj900">sj900</a> for integrating and testing the `WeightedPDBSampler` within the `PDBDataset` and for adding initial support for MSA and template parsing!
 
 - <a href="https://github.com/xluo233">@xluo233</a> again for contributing the logic for computing the model selection score as well as the unresolved rasa!
 
@@ -69,7 +69,7 @@ from alphafold3_pytorch.utils.model_utils import exclusive_cumsum
 
 alphafold3 = Alphafold3(
     dim_atom_inputs = 77,
-    dim_template_feats = 44
+    dim_template_feats = 108
 )
 
 # mock inputs
@@ -91,7 +91,7 @@ is_molecule_types = torch.randint(0, 2, (2, seq_len, 5)).bool()
 is_molecule_mod = torch.randint(0, 2, (2, seq_len, 4)).bool()
 molecule_ids = torch.randint(0, 32, (2, seq_len))
 
-template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
+template_feats = torch.randn(2, 2, seq_len, seq_len, 108)
 template_mask = torch.ones((2, 2)).bool()
 
 msa = torch.randn(2, 7, seq_len, 32)
@@ -197,7 +197,7 @@ alphafold3 = Alphafold3(
     dim_atom_inputs = 3,
     dim_atompair_inputs = 5,
     atoms_per_window = 27,
-    dim_template_feats = 44,
+    dim_template_feats = 108,
     num_dist_bins = 38,
     num_molecule_mods = 0,
     confidence_head_kwargs = dict(

alphafold3_pytorch/__init__.py

@@ -9,8 +9,6 @@
     SmoothLDDTLoss,
     WeightedRigidAlign,
     MultiChainPermutationAlignment,
-    ExpressCoordinatesInFrame,
-    RigidFrom3Points,
     ComputeAlignmentError,
     CentreRandomAugmentation,
     TemplateEmbedder,
@@ -70,6 +68,10 @@
     create_trainer_from_yaml,
     create_trainer_from_conductor_yaml
 )
+from alphafold3_pytorch.utils.model_utils import (
+    ExpressCoordinatesInFrame,
+    RigidFrom3Points,
+)
 
 __all__ = [
     Attention,

alphafold3_pytorch/alphafold3.py

@@ -70,6 +70,8 @@
 )
 
 from alphafold3_pytorch.utils.model_utils import (
+    ExpressCoordinatesInFrame,
+    RigidFrom3Points,
     calculate_weighted_rigid_align_weights,
 )
 
@@ -3899,93 +3901,6 @@ def forward(
         permuted_true_coords = labels["true_coords"].detach()
         return permuted_true_coords
 
-class ExpressCoordinatesInFrame(Module):
-    """ Algorithm  29 """
-
-    def __init__(
-        self,
-        eps = 1e-8
-    ):
-        super().__init__()
-        self.eps = eps
-
-    @typecheck
-    def forward(
-        self,
-        coords: Float['b m 3'],
-        frame: Float['b m 3 3'] | Float['b 3 3'] | Float['3 3']
-    ) -> Float['b m 3']:
-        """
-        coords: coordinates to be expressed in the given frame
-        frame: frame defined by three points
-        """
-
-        if frame.ndim == 2:
-            frame = rearrange(frame, 'fr fc -> 1 1 fr fc')
-        elif frame.ndim == 3:
-            frame = rearrange(frame, 'b fr fc -> b 1 fr fc')
-
-        # Extract frame atoms
-        a, b, c = frame.unbind(dim=-1)
-        w1 = l2norm(a - b, eps=self.eps)
-        w2 = l2norm(c - b, eps=self.eps)
-
-        # Build orthonormal basis
-        e1 = l2norm(w1 + w2, eps=self.eps)
-        e2 = l2norm(w2 - w1, eps=self.eps)
-        e3 = torch.cross(e1, e2, dim=-1)
-
-        # Project onto frame basis
-        d = coords - b
-
-        transformed_coords = torch.stack((
-            einsum(d, e1, '... i, ... i -> ...'),
-            einsum(d, e2, '... i, ... i -> ...'),
-            einsum(d, e3, '... i, ... i -> ...'),
-        ), dim=-1)
-
-        return transformed_coords
-
-class RigidFrom3Points(Module):
-    """
-    Algorithm 21 in Section 1.8.1 in Alphafold2 paper
-    https://www.nature.com/articles/s41586-021-03819-2
-    """
-
-    @typecheck
-    def forward(
-        self,
-        three_points: Tuple[Float['... 3'], Float['... 3'], Float['... 3']] | Float['3 ... 3']
-    ) -> Tuple[Float['... 3 3'], Float['... 3']]:
-
-        if isinstance(three_points, tuple):
-            three_points = torch.stack(three_points)
-
-        # allow for any number of leading dimensions
-
-        (x1, x2, x3), unpack_one = pack_one(three_points, 'three * d')
-
-        # main algorithm
-
-        v1 = x3 - x2
-        v2 = x1 - x2
-
-        e1 = l2norm(v1)
-        u2 = v2 - e1 @ (e1.t() @ v2)
-        e2 = l2norm(u2)
-
-        e3 = torch.cross(e1, e2, dim = -1)
-
-        R = torch.stack((e1, e2, e3), dim = -1)
-        t = x2
-
-        # unpack
-
-        R = unpack_one(R, '* r1 r2')
-        t = unpack_one(t, '* c')
-
-        return R, t
-
 class ComputeAlignmentError(Module):
     """ Algorithm 30 """
 
@@ -6364,7 +6279,16 @@ def forward(
         if hard_debug:
             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')
+
+            is_biomolecule = ~(
+                (~is_molecule_types[..., IS_BIOMOLECULE_INDICES].any(dim=-1))
+                | (exists(is_molecule_mod) and is_molecule_mod.any(dim=-1))
+            )
+            maybe(hard_validate_atom_indices_ascending)(
+                atom_indices_for_frame,
+                'atom_indices_for_frame',
+                mask=is_biomolecule,
+            )
 
         # soft validate
 
@@ -6505,11 +6429,6 @@ def forward(
         mask = molecule_atom_lens > 0
         pairwise_mask = to_pairwise_mask(mask)
 
-        # prepare mask for msa module and template embedder
-        # which is equivalent to the `is_protein` of the `is_molecular_types` input
-
-        is_protein_mask = is_molecule_types[..., IS_PROTEIN_INDEX]
-
         # init recycled single and pairwise
 
         detach_when_recycling = default(detach_when_recycling, self.detach_when_recycling)
@@ -6546,7 +6465,6 @@ def forward(
                     templates = templates,
                     template_mask = template_mask,
                     pairwise_repr = pairwise,
-                    mask = is_protein_mask
                 )
 
                 pairwise = embedded_template + pairwise
@@ -6558,7 +6476,6 @@ def forward(
                     msa = msa,
                     single_repr = single,
                     pairwise_repr = pairwise,
-                    mask = is_protein_mask,
                     msa_mask = msa_mask,
                     additional_msa_feats = additional_msa_feats
                 )
@@ -6961,12 +6878,9 @@ def forward(
                 pred_frames, _ = self.rigid_from_three_points(pred_three_atoms)
 
                 # determine mask
-                # must be residue or nucleotide with greater than 0 atoms
+                # must be amino acid, nucleotide, or ligand with greater than 0 atoms
 
-                align_error_mask = (
-                    is_molecule_types[..., IS_BIOMOLECULE_INDICES].any(dim=-1)
-                    & valid_atom_indices_for_frame
-                )
+                align_error_mask = valid_atom_indices_for_frame
 
                 # align error
 
@@ -6982,7 +6896,7 @@ def forward(
 
                 pae_labels = distance_to_bins(align_error, self.pae_bins)
 
-                # set ignore index for invalid molecules or frames (TODO: figure out what is meant by invalid frame)
+                # set ignore index for invalid molecules or frames
 
                 pair_align_error_mask = to_pairwise_mask(align_error_mask)