quick fix (#133)

* hack to fix pytest-split-tests

* try again

* missing comma

* some lint

* some cleanup

* skip rasa test if dssp not installed

* another fix

* attempt to install dssp in test env

* sudo

* skip rasa test for now

Author: lucidrains

README.md

@@ -42,7 +42,7 @@ A fork with full Lightning + Hydra support is being maintained by <a href="https
 
 - <a href="https://github.com/sj900">sj900</a> for integrating and testing the `WeightedPDBSampler` within the `PDBDataset`!
 
-- <a href="https://github.com/xluo233">@xluo233</a> again for contributing the logic for computing the model selection score!
+- <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!
 
 - <a href="https://github.com/wufandi">Fandi</a> for discovering a few inconsistencies in the elucidated atom diffusion module with the supplementary
 

alphafold3_pytorch/alphafold3.py

@@ -1995,11 +1995,9 @@ def forward(
         mask: Bool['b n'] | None = None,
         windowed_mask: Bool['b nw w (w*2)'] | None = None
     ):
-        w = self.attn_window_size
+        w, serial = self.attn_window_size, self.serial
         has_windows = exists(w)
 
-        serial = self.serial
-
         # handle windowing
 
         pairwise_is_windowed = pairwise_repr.ndim == 5
@@ -2022,9 +2020,9 @@ def forward(
 
         # main transformer
 
-        if self.serial and should_checkpoint(self, (noised_repr, single_repr, pairwise_repr)):
+        if serial and should_checkpoint(self, (noised_repr, single_repr, pairwise_repr)):
             to_layers_fn = self.to_checkpointed_serial_layers
-        elif self.serial:
+        elif serial:
             to_layers_fn = self.to_serial_layers
         else:
             to_layers_fn = self.to_parallel_layers
@@ -4067,7 +4065,7 @@ def __init__(
         other_cutoff: float = 15.0,
         contact_mask_threshold: float = 8.0,
         is_fine_tuning: bool = False,
-        weight_dict_config: dict = None
+        weight_dict_config: dict = None,
         dssp_path: str = 'mkdssp',
     ):
 
@@ -4105,7 +4103,7 @@ def compute_gpde(
         dist_probs = F.softmax(dist_logits, dim=-1)
 
         # for distances greater than the last breaks
-        dist_breaks = F.pdb(dist_breaks, (0, 1), value = 1e6)
+        dist_breaks = F.pad(dist_breaks, (0, 1), value = 1e6)
         contact_mask = dist_breaks < self.contact_mask_threshold
 
         contact_prob = einx.where(
@@ -4389,7 +4387,7 @@ def compute_unresolved_rasa(
         molecule_atom_lens: Int['b n'],
         atom_pos: Float['b m 3'], 
         atom_mask: Bool['b m'],
-    ) -> Float['b']:
+    ) -> Float[' b']:
 
         unresolved_rasa = [self._compute_unresolved_rasa(*args) for args in 
                            zip(unresolved_cid, unresolved_residue_mask, asym_id, molecule_ids, molecule_atom_lens, atom_pos, atom_mask)]

alphafold3_pytorch/inputs.py

@@ -316,7 +316,7 @@ def __getitem__(self, idx: int) -> AtomInput:
 @typecheck
 def atom_ref_pos_to_atompair_inputs(
     atom_ref_pos: Float['m 3'],
-    atom_ref_space_uid: Int['m'] | None = None,
+    atom_ref_space_uid: Int[' m'] | None = None,
 ) -> Float['m m 5']:
 
     # Algorithm 5 - lines 2-6
@@ -691,7 +691,7 @@ class MoleculeLengthMoleculeInput:
     src_tgt_atom_indices:       Int['n 2']
     token_bonds:                Bool['n n'] | None = None
     one_token_per_atom:         List[bool] | None = None
-    is_molecule_mod:            Bool['n num_mods'] | Bool['n'] | None = None
+    is_molecule_mod:            Bool['n num_mods'] | Bool[' n'] | None = None
     molecule_atom_indices:      List[int | None] | None = None
     distogram_atom_indices:     List[int | None] | None = None
     missing_atom_indices:       List[Int[' _'] | None] | None = None

alphafold3_pytorch/trainer.py

@@ -35,6 +35,7 @@
 
 import torch
 from torch import Tensor
+from torch.nn import Module
 from torch.optim import Adam, Optimizer
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import Sampler, Dataset, DataLoader as OrigDataLoader

pyproject.toml

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

tests/test_af3.py

@@ -4,10 +4,13 @@
 import pytest
 import random
 import itertools
+import subprocess
 from pathlib import Path
 
 import torch
 
+from collections import namedtuple
+
 from alphafold3_pytorch import (
     SmoothLDDTLoss,
     WeightedRigidAlign,
@@ -30,6 +33,7 @@
     ConfidenceHeadLogits,
     ComputeModelSelectionScore,
     ComputeModelSelectionScore,
+    collate_inputs_to_batched_atom_input
 )
 
 from alphafold3_pytorch.configs import (
@@ -40,12 +44,20 @@
 from alphafold3_pytorch.alphafold3 import (
     mean_pool_with_lens,
     repeat_consecutive_with_lens,
-    full_pairwise_repr_to_windowed
+    full_pairwise_repr_to_windowed,
+    get_cid_molecule_type,
 )
 
 from alphafold3_pytorch.inputs import (
     IS_MOLECULE_TYPES,
-    atom_ref_pos_to_atompair_inputs
+    IS_PROTEIN,
+    atom_ref_pos_to_atompair_inputs,
+    molecule_to_atom_input,
+    pdb_input_to_molecule_input,
+    PDBInput,
+    PDBDataset,
+    default_extract_atom_feats_fn,
+    default_extract_atompair_feats_fn
 )
 
 def test_atom_ref_pos_to_atompair_inputs():
@@ -1090,27 +1102,14 @@ def test_model_selection_score():
 
 def test_unresolved_protein_rasa():
 
-    from collections import namedtuple
+    # skip the test if dssp not installed
 
-    from alphafold3_pytorch.inputs import (
-        IS_MOLECULE_TYPES, 
-        PDBInput,
-        default_extract_atom_feats_fn,
-        default_extract_atompair_feats_fn
-        
-    )
+    try:
+        subprocess.check_output(["which", "mkdssp"])
+    except:
+        pytest.skip("mkdssp not found, test_unresolved_protein_rasa skipped")
 
-    from alphafold3_pytorch.inputs import (
-        PDBDataset,
-        molecule_to_atom_input,
-        pdb_input_to_molecule_input,
-        IS_PROTEIN,
-    )
-
-    from alphafold3_pytorch import collate_inputs_to_batched_atom_input
-    from alphafold3_pytorch.alphafold3 import (
-        get_cid_molecule_type,
-    )
+    # rest of the test
 
     mmcif_filepath = os.path.join('data', 'test', '7a4d-assembly1.cif')
     pdb_input = PDBInput(mmcif_filepath)