Add full constraints set (#268)

* Update inputs.py

* Update test_dataloading.py

* Update inputs.py

* Update weighted_pdb_sampler.py

* Update inputs.py

* Update plm.py

* Update alphafold3.py

* Update template_parsing.py

* Update inputs.py

* Update inputs.py

* Update alphafold3.py

* Update alphafold3.py

Author: amorehead

alphafold3_pytorch/alphafold3.py

@@ -54,13 +54,17 @@
 )
 
 from alphafold3_pytorch.inputs import (
+    CONSTRAINT_DIMS,
+    CONSTRAINTS,
+    CONSTRAINTS_MASK_VALUE,
     IS_MOLECULE_TYPES,
     IS_PROTEIN_INDEX,
     IS_DNA_INDEX,
     IS_RNA_INDEX,
     IS_LIGAND_INDEX,
     IS_METAL_ION_INDEX,
     IS_BIOMOLECULE_INDICES,
+    IS_NON_PROTEIN_INDICES,
     IS_PROTEIN,
     IS_DNA,
     IS_RNA,
@@ -5954,7 +5958,7 @@ def __init__(
         pdb_training_set=True,
         plm_embeddings: PLMEmbedding | tuple[PLMEmbedding, ...] | None = None,
         plm_kwargs: dict | tuple[dict, ...] | None = None,
-        constraint_embeddings: int | None = None,
+        constraints: List[CONSTRAINTS] | None = None,
     ):
         super().__init__()
 
@@ -5983,10 +5987,18 @@ def __init__(
 
         # optional pairwise token constraint embeddings
 
-        self.constraint_embeddings = constraint_embeddings
+        self.constraints = constraints
 
-        if exists(constraint_embeddings):
-            self.constraint_embeds = LinearNoBias(constraint_embeddings, dim_pairwise)
+        if exists(constraints):
+            self.constraint_embeds = nn.ModuleList(
+                [
+                    LinearNoBias(CONSTRAINT_DIMS[constraint], dim_pairwise)
+                    for constraint in constraints
+                ]
+            )
+            self.learnable_constraint_masks = nn.ParameterList(
+                [nn.Parameter(torch.randn(1)) for _ in constraints]
+            )
 
         # residue or nucleotide modifications
 
@@ -6538,21 +6550,30 @@ def forward(
 
         # handle maybe pairwise token constraint embeddings
 
-        if exists(self.constraint_embeddings):
+        if exists(self.constraints):
             assert exists(
                 token_constraints
             ), "`token_constraints` must be provided to use constraint embeddings."
 
-            pairwise_constraint_embeds = self.constraint_embeds(token_constraints)
-            pairwise_init = pairwise_init + pairwise_constraint_embeds
+            for i, constraint in enumerate(self.constraints):
+                constraint_slice = slice(i, i + CONSTRAINT_DIMS[constraint])
+
+                token_constraint = torch.where(
+                    # replace fixed constraint mask values with learnable mask
+                    token_constraints[..., constraint_slice] == CONSTRAINTS_MASK_VALUE,
+                    self.learnable_constraint_masks[i],
+                    token_constraints[..., constraint_slice],
+                )
+
+                pairwise_init = pairwise_init + self.constraint_embeds[i](token_constraint)
 
         # handle maybe protein language model (PLM) embeddings
 
         if exists(self.plms):
             molecule_aa_ids = torch.where(
-                molecule_ids < 0,
-                NUM_HUMAN_AMINO_ACIDS,
-                molecule_ids.clamp(max=NUM_HUMAN_AMINO_ACIDS),
+                is_molecule_types[..., IS_NON_PROTEIN_INDICES].any(dim=-1),
+                -1,
+                molecule_ids,
             )
 
             plm_embeds = [plm(molecule_aa_ids) for plm in self.plms]

alphafold3_pytorch/data/template_parsing.py

@@ -121,13 +121,8 @@ def parse_m8(
             template_release_date = extract_mmcif_metadata_field(
                 template_mmcif_object, "release_date"
             )
-            if not (
-                exists(template_cutoff_date)
-                and datetime.strptime(template_release_date, "%Y-%m-%d") <= template_cutoff_date
-            ):
+            if exists(template_cutoff_date) and datetime.strptime(template_release_date, "%Y-%m-%d") > template_cutoff_date:
                 continue
-            elif not_exists(template_cutoff_date):
-                pass
             template_biomol = _from_mmcif_object(
                 template_mmcif_object, chain_ids=set(template_chain)
             )

alphafold3_pytorch/data/weighted_pdb_sampler.py

@@ -191,7 +191,7 @@ def __init__(
         alpha_prot: float = 3.0,
         alpha_nuc: float = 3.0,
         alpha_ligand: float = 1.0,
-        pdb_ids_to_skip: List[str] = [],
+        pdb_ids_to_skip: List[str] | None = None,
         pdb_ids_to_keep: list[str] | None = None,
     ):
         # Load chain and interface mappings
@@ -210,7 +210,7 @@ def __init__(
         interface_mapping = pl.read_csv(interface_mapping_path)
 
         # Filter out unwanted PDB IDs
-        if len(pdb_ids_to_skip) > 0:
+        if exists(pdb_ids_to_skip) and len(pdb_ids_to_skip) > 0:
             chain_mapping = chain_mapping.filter(pl.col("pdb_id").is_in(pdb_ids_to_skip).not_())
             interface_mapping = interface_mapping.filter(
                 pl.col("pdb_id").is_in(pdb_ids_to_skip).not_()