fix types

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -3492,12 +3492,14 @@ def compute_modified_residue_score(
         return plddt_mean
 
 # model selection
+
+@typecheck
 def get_cid_molecule_type(
     cid: int,
     asym_id: Int[' n'],
-    is_molecule_types: Bool['n {IS_MOLECULE_TYPES}'],
+    is_molecule_types: Bool[f'n {IS_MOLECULE_TYPES}'],
     return_one_hot: bool = False,
-    ) -> int | Bool[' {IS_MOLECULE_TYPES}']:
+) -> int | Bool[f' {IS_MOLECULE_TYPES}']:
     """
     
     get the molecule type for where asym_id == cid
@@ -3517,6 +3519,7 @@ def get_cid_molecule_type(
 
 class ComputeModelSelectionScore(Module):
 
+    @typecheck
     def __init__(
         self,
         eps: float = 1e-8,
@@ -3535,6 +3538,7 @@ def __init__(
 
         self.register_buffer('dist_breaks', dist_breaks)
 
+    @typecheck
     def compute_gpde(
         self,
         pde_logits: Float['b pde n n'],
@@ -3570,6 +3574,7 @@ def compute_gpde(
 
         return gpde
 
+    @typecheck
     def compute_lddt(
         self,
         pred_coords: Float['b m 3'],
@@ -3628,13 +3633,14 @@ def compute_lddt(
 
         return lddt_mean
 
+    @typecheck
     def compute_chain_pair_lddt(
         self,
         asym_mask_a: Bool['b m'] | Bool [' m'],
         asym_mask_b: Bool['b m'] | Bool [' m'],
         pred_coords: Float['b m 3'] | Float['m 3'],
         true_coords: Float['b m 3'] | Float['m 3'], 
-        is_molecule_types: Int['b m {IS_MOLECULE_TYPES}'] | Int['m {IS_MOLECULE_TYPES}'],
+        is_molecule_types: Bool[f'b m {IS_MOLECULE_TYPES}'] | Bool[f'm {IS_MOLECULE_TYPES}'],
         coords_mask: Bool['b m'] | Bool [' m'] | None = None,
     ) -> Float[' b']:
         """
@@ -3665,6 +3671,7 @@ def compute_chain_pair_lddt(
 
         return lddt
 
+    @typecheck
     def get_lddt_weight(
         self,
         type_chain_a,
@@ -3718,7 +3725,7 @@ def get_lddt_weight(
         weight_dict = fine_tuning_dict if is_fine_tuning else initial_training_dict
 
         if lddt_type == 'unresolved':
-            weight =  weight_dict.get(lddt_type, None).get(lddt_type, None)
+            weight =  weight_dict.get(lddt_type, {}).get(lddt_type, None)
             assert weight
             return weight
 
@@ -3728,6 +3735,7 @@ def get_lddt_weight(
         assert weight, f"Weight not found for {interface_type} {lddt_type}"
         return weight
 
+    @typecheck
     def compute_weighted_lddt(
         self,
         # atom level input
@@ -3736,7 +3744,7 @@ def compute_weighted_lddt(
         atom_mask: Bool['b m'] | None,
         # token level input
         asym_id: Int['b n'],
-        is_molecule_types: Bool['b n {IS_MOLECULE_TYPES}'],
+        is_molecule_types: Bool[f'b n {IS_MOLECULE_TYPES}'],
         molecule_atom_lens: Int['b n'],
         # additional input
         chains_list: List[Tuple[int, int] | Tuple[int]],

pyproject.toml

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