able to force sampling or return of loss, no matter the state of the labels being passed in

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -3161,7 +3161,7 @@ def forward(
         plddt_labels: Int['b n'] | None = None,
         resolved_labels: Int['b n'] | None = None,
         return_loss_breakdown = False,
-        return_loss_if_possible: bool = True,
+        return_loss: bool = None,
         num_rollout_steps: int = 20,
         rollout_show_tqdm_pbar: bool = False
     ) -> Float['b m 3'] | Float[''] | Tuple[Float[''], LossBreakdown]:
@@ -3317,11 +3317,15 @@ def forward(
 
         has_labels = any([*map(exists, all_labels)])
 
-        return_loss = atom_pos_given or has_labels
+        can_return_loss = atom_pos_given or has_labels
+
+        # default whether to return loss by whether labels or atom positions are given
+
+        return_loss = default(return_loss, can_return_loss)
 
         # if neither atom positions or any labels are passed in, sample a structure and return
 
-        if not return_loss_if_possible or not return_loss:
+        if not return_loss:
             return self.edm.sample(
                 num_sample_steps = num_sample_steps,
                 atom_feats = atom_feats,
@@ -3335,6 +3339,16 @@ def forward(
                 molecule_atom_lens = molecule_atom_lens
             )
 
+        # if being forced to return loss, but do not have sufficient information to return losses, just return 0
+
+        if return_loss and not can_return_loss:
+            zero = self.zero.requires_grad_()
+
+            if not return_loss_breakdown:
+                return zero
+
+            return zero, LossBreakdown(*((zero,) * 11))
+
         # losses default to 0
 
         distogram_loss = diffusion_loss = confidence_loss = pae_loss = pde_loss = plddt_loss = resolved_loss = self.zero

pyproject.toml

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

tests/test_af3.py

@@ -37,9 +37,6 @@
     atom_ref_pos_to_atompair_inputs
 )
 
-def join(str, delimiter = ','):
-    return delimiter.join(str)
-
 def test_atom_ref_pos_to_atompair_inputs():
     atom_ref_pos = torch.randn(16, 3)
     atom_ref_space_uid = torch.ones(16).long()
@@ -409,14 +406,8 @@ def test_distogram_head():
 
     logits = distogram_head(pairwise_repr)
 
-@pytest.mark.parametrize(
-    join([
-        'window_atompair_inputs',
-        'stochastic_frame_average'
-    ]), [
-        (True, False),
-        (True, False)
-    ])
+@pytest.mark.parametrize('window_atompair_inputs', (True, False))
+@pytest.mark.parametrize('stochastic_frame_average', (True, False))
 def test_alphafold3(
     window_atompair_inputs: bool,
     stochastic_frame_average: bool
@@ -572,6 +563,78 @@ def test_alphafold3_without_msa_and_templates():
 
     loss.backward()
 
+def test_alphafold3_force_return_loss():
+    seq_len = 16
+    molecule_atom_lens = torch.randint(1, 3, (2, seq_len))
+    atom_seq_len = molecule_atom_lens.sum(dim = -1).amax()
+
+    atom_inputs = torch.randn(2, atom_seq_len, 77)
+    atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
+    additional_molecule_feats = torch.randn(2, seq_len, 10)
+
+    atom_pos = torch.randn(2, atom_seq_len, 3)
+    molecule_atom_indices = molecule_atom_lens - 1
+
+    distance_labels = torch.randint(0, 38, (2, seq_len, seq_len))
+    pae_labels = torch.randint(0, 64, (2, seq_len, seq_len))
+    pde_labels = torch.randint(0, 64, (2, seq_len, seq_len))
+    plddt_labels = torch.randint(0, 50, (2, seq_len))
+    resolved_labels = torch.randint(0, 2, (2, seq_len))
+
+    alphafold3 = Alphafold3(
+        dim_atom_inputs = 77,
+        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
+        ),
+        pairformer_stack = dict(
+            depth = 2
+        ),
+        diffusion_module_kwargs = dict(
+            atom_encoder_depth = 1,
+            token_transformer_depth = 1,
+            atom_decoder_depth = 1,
+        ),
+    )
+
+    sampled_atom_pos = alphafold3(
+        num_recycling_steps = 2,
+        atom_inputs = atom_inputs,
+        molecule_atom_lens = molecule_atom_lens,
+        atompair_inputs = atompair_inputs,
+        additional_molecule_feats = additional_molecule_feats,
+        atom_pos = atom_pos,
+        molecule_atom_indices = molecule_atom_indices,
+        distance_labels = distance_labels,
+        pae_labels = pae_labels,
+        pde_labels = pde_labels,
+        plddt_labels = plddt_labels,
+        resolved_labels = resolved_labels,
+        return_loss_breakdown = True,
+        return_loss = False # force sampling even if labels are given
+    )
+
+    assert sampled_atom_pos.ndim == 3
+
+    loss, _ = alphafold3(
+        num_recycling_steps = 2,
+        atom_inputs = atom_inputs,
+        molecule_atom_lens = molecule_atom_lens,
+        atompair_inputs = atompair_inputs,
+        additional_molecule_feats = additional_molecule_feats,
+        return_loss_breakdown = True,
+        return_loss = True # force returning loss even if no labels given
+    )
+
+    assert loss == 0.
+
 # test creation from config
 
 def test_alphafold3_config():