add end to end test for model selection (#177)

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -157,8 +157,6 @@
 
 # constants
 
-SCORED_SAMPLE = Tuple[int, Float["b m 3"], Float[" b"], Float[" b"]] # type: ignore
-
 LinearNoBias = partial(Linear, bias = False)
 
 # always use non reentrant checkpointing
@@ -3470,6 +3468,13 @@ class ConfidenceHeadLogits(NamedTuple):
     plddt: Float['b plddt m']
     resolved: Float['b 2 m']
 
+class Alphafold3Logits(NamedTuple):
+    pae: Float['b pae n n'] |  None
+    pde: Float['b pde n n']
+    plddt: Float['b plddt m']
+    resolved: Float['b 2 m']
+    distance: Float['b dist m m'] | Float['b dist n n'] | None
+
 class ConfidenceHead(Module):
     """ Algorithm 31 """
 
@@ -4281,6 +4286,13 @@ def _protein_structure_from_feature(
 
     return builder.get_structure()
 
+ScoredSample = Tuple[int, Float["b m 3"], Float[" b"], Float[" b"]] # type: ignore
+
+class ScoreDetails(NamedTuple):
+    best_gpde_index: int
+    best_lddt_index: int
+    score: Float[' b']
+    scored_samples: ScoredSample
 
 class ComputeModelSelectionScore(Module):
     """Compute model selection score."""
@@ -4787,15 +4799,20 @@ def compute_unresolved_rasa(
     def compute_model_selection_score(
         self,
         batch: BatchedAtomInput,
-        samples: List[Tuple[Float["b m 3"], Float["b pde n n"], Float["b dist n n"]]],  
+        samples: List[Tuple[
+            Float["b m 3"],
+            Float["b pde n n"],
+            Float["b dist n n"]
+        ]],
         is_fine_tuning: bool = None,
-        return_top_model: bool = False,
+        return_details: bool = False,
         return_unweighted_scores: bool = False,
         compute_rasa: bool = False,
         unresolved_cid: List[int] | None = None,
         unresolved_residue_mask: Bool["b n"] | None = None,  
         missing_chain_index: int = -1,
-    ) -> Float[" b"] | Tuple[Float[" b"], SCORED_SAMPLE]:  
+    ) -> Float[" b"] | ScoreDetails:
+
         """Compute the model selection score for an input batch and corresponding (sampled) atom
         positions.
 
@@ -4841,7 +4858,7 @@ def compute_model_selection_score(
 
         # score samples
 
-        scored_samples: List[SCORED_SAMPLE] = []
+        scored_samples: List[ScoredSample] = []
 
         for sample_idx, sample in enumerate(samples):
             atom_pos_pred, pde_logits, dist_logits = sample
@@ -4871,19 +4888,73 @@ def compute_model_selection_score(
 
             scored_samples.append((sample_idx, atom_pos_pred, weighted_lddt, gpde))
 
-        top_ranked_sample = max(
-            scored_samples, key=lambda x: x[-1].mean()
-        )  # rank by batch-averaged gPDE
-        best_of_5_sample = max(
-            scored_samples, key=lambda x: x[-2].mean()
-        )  # rank by batch-averaged lDDT
+        # quick collate
+
+        *_, all_weighted_lddt, all_gpde = zip(*scored_samples)
+
+        # rank by batch-averaged gPDE
 
-        model_selection_score = (top_ranked_sample[-2] + best_of_5_sample[-2]) / 2
+        best_gpde_index = torch.stack(all_gpde).mean(dim = -1).topk(1).indices.item()
 
-        if return_top_model:
-            return model_selection_score, top_ranked_sample
+        # rank by batch-averaged lDDT
 
-        return model_selection_score
+        best_lddt_index = torch.stack(all_weighted_lddt).mean(dim = -1).topk(1).indices.item()
+
+        # some weighted score
+
+        model_selection_score = (
+            scored_samples[best_gpde_index][-2] +
+            scored_samples[best_lddt_index][-2]
+        ) / 2
+
+        if not return_details:
+            return model_selection_score
+
+        score_details = ScoreDetails(
+            best_gpde_index = best_gpde_index,
+            best_lddt_index = best_lddt_index,
+            score = model_selection_score,
+            scored_samples = scored_samples
+        )
+
+        return score_details
+
+    @typecheck
+    def forward(
+        self,
+        alphafolds: Tuple[Alphafold3],
+        batched_atom_inputs: BatchedAtomInput,
+        **kwargs
+    ) -> Float[" b"] | ScoreDetails:
+
+        """
+        give this a tuple of all the Alphafolds and a batch of atomic inputs
+        it will select the best one by the model selection score by returning the index of the Tuple
+        """
+
+        samples = []
+
+        with torch.no_grad():
+            for alphafold in alphafolds:
+                alphafold.eval()
+
+                pred_atom_pos, logits = alphafold(
+                    **batched_atom_inputs.model_forward_dict(),
+                    return_loss = False,
+                    return_confidence_head_logits = True,
+                    return_distogram_head_logits = True
+                )
+
+                samples.append((pred_atom_pos, logits.pde, logits.distance))
+
+
+        scores = self.compute_model_selection_score(
+            batched_atom_inputs,
+            samples = samples,
+            **kwargs
+        )
+
+        return scores
 
 # main class
 
@@ -5383,8 +5454,7 @@ def forward(
         hard_validate: bool = False
     ) -> (
         Float['b m 3'] |
-        Tuple[Float['b m 3'], ConfidenceHeadLogits] |
-        Tuple[Float['b m 3'], ConfidenceHeadLogits, Float['b l n n'] | Float['b l m m']] |
+        Tuple[Float['b m 3'], ConfidenceHeadLogits | Alphafold3Logits] |
         Float[''] |
         Tuple[Float[''], LossBreakdown]
     ):
@@ -5673,16 +5743,17 @@ def forward(
                 return_pae_logits = True
             )
 
-            if not return_distogram_head_logits:
-                return sampled_atom_pos, confidence_head_logits
+            returned_logits = confidence_head_logits
 
-            distogram_head_logits = self.distogram_head(pairwise.clone().detach())
+            if return_distogram_head_logits:
+                distogram_head_logits = self.distogram_head(pairwise.clone().detach())
 
-            return (
-                sampled_atom_pos,
-                confidence_head_logits,
-                distogram_head_logits,
-            )
+                returned_logits = Alphafold3Logits(
+                    **confidence_head_logits._asdict(),
+                    distance = distogram_head_logits
+                )
+
+            return sampled_atom_pos, returned_logits
 
         # if being forced to return loss, but do not have sufficient information to return losses, just return 0
 

pyproject.toml

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

tests/test_af3.py

@@ -39,6 +39,8 @@
     collate_inputs_to_batched_atom_input
 )
 
+from alphafold3_pytorch.mocks import MockAtomDataset
+
 from alphafold3_pytorch.configs import (
     Alphafold3Config,
     create_alphafold3_from_yaml
@@ -61,7 +63,7 @@
     PDBInput,
     PDBDataset,
     default_extract_atom_feats_fn,
-    default_extract_atompair_feats_fn
+    default_extract_atompair_feats_fn,
 )
 
 from alphafold3_pytorch.utils.model_utils import exclusive_cumsum
@@ -1098,7 +1100,6 @@ def test_model_selection_score():
         for chain_len in chain_length
     ]).long()
 
-
     is_molecule_types = torch.zeros_like(asym_id)
     is_molecule_types = torch.nn.functional.one_hot(is_molecule_types, 5).bool()
 
@@ -1122,6 +1123,71 @@ def test_model_selection_score():
         is_fine_tuning=False
     )
 
+def test_model_selection_score_end_to_end():
+
+    # prepare two atom inputs for evaluating model selection
+
+    mock_atom_dataset = MockAtomDataset(10)
+
+    atom_inputs = [mock_atom_dataset[0], mock_atom_dataset[1]]
+    batched_atom_input = collate_inputs_to_batched_atom_input(atom_inputs, atoms_per_window=27)
+
+    # two models to be selected
+
+    alphafold3_kwargs = dict(
+        dim_atom_inputs = 77,
+        dim_pairwise = 8,
+        dim_single = 8,
+        dim_token = 8,
+        atoms_per_window = 27,
+        dim_template_feats = 44,
+        num_dist_bins = 38,
+        confidence_head_kwargs = dict(
+            pairformer_depth = 1
+        ),
+        template_embedder_kwargs = dict(
+            pairformer_stack_depth = 1
+        ),
+        msa_module_kwargs = dict(
+            depth = 1,
+            dim_msa = 8,
+        ),
+        pairformer_stack=dict(
+            depth=1,
+            pair_bias_attn_dim_head = 4,
+            pair_bias_attn_heads = 2,
+        ),
+        diffusion_module_kwargs=dict(
+            atom_encoder_depth=1,
+            token_transformer_depth=1,
+            atom_decoder_depth=1,
+            atom_decoder_kwargs = dict(
+                attn_pair_bias_kwargs = dict(
+                    dim_head = 4
+                )
+            ),
+            atom_encoder_kwargs = dict(
+                attn_pair_bias_kwargs = dict(
+                    dim_head = 4
+                )
+            )
+        ),
+    )
+
+    alphafold3_one = Alphafold3(**alphafold3_kwargs)
+    alphafold3_two = Alphafold3(**alphafold3_kwargs)
+
+    alphafolds = (alphafold3_one, alphafold3_two)
+
+    # evaluate
+
+    compute_model_selection_score = ComputeModelSelectionScore()
+
+    details = compute_model_selection_score(alphafolds, batched_atom_input, return_details = True)
+
+    best_alphafold_by_lddt = alphafolds[details.best_lddt_index]
+    assert isinstance(best_alphafold_by_lddt, Alphafold3)
+
 def test_unresolved_protein_rasa():
 
     # rest of the test