Fix various edge case bugs raised during training-time cropping (#149)

* Update inputs.py

* Update biomolecule.py

Author: amorehead

alphafold3_pytorch/common/biomolecule.py

@@ -399,6 +399,8 @@ def spatial_crop(
             token_center_atom_mask[self.chain_id == chain_1] = True
         elif exists(chain_2):
             token_center_atom_mask[self.chain_id == chain_2] = True
+        else:
+            raise ValueError("At least one chain ID must be specified for spatial cropping.")
 
         # potentially filter candidate token center atoms by interface proximity
 
@@ -452,7 +454,16 @@ def crop(
     ) -> "Biomolecule":
         """Crop a Biomolecule using a randomly-sampled cropping function."""
         n_res = min(n_res, len(self.atom_mask))
-        crop_fn_weights = [contiguous_weight, spatial_weight, spatial_interface_weight]
+        if exists(chain_1) and exists(chain_2):
+            crop_fn_weights = [contiguous_weight, spatial_weight, spatial_interface_weight]
+        elif exists(chain_1) or exists(chain_2):
+            crop_fn_weights = [contiguous_weight, spatial_weight + spatial_interface_weight, 0.0]
+        else:
+            crop_fn_weights = [
+                contiguous_weight + spatial_weight + spatial_interface_weight,
+                0.0,
+                0.0,
+            ]
         crop_fns = [
             partial(self.contiguous_crop, n_res=n_res),
             partial(

alphafold3_pytorch/inputs.py

@@ -2113,15 +2113,16 @@ def pdb_input_to_molecule_input(
     """Convert a PDBInput to a MoleculeInput."""
     i = pdb_input
 
+    filepath = pdb_input.mmcif_filepath
+    file_id = os.path.splitext(os.path.basename(filepath))[0]
+
     # acquire a `Biomolecule` object for the given `PDBInput`
 
     if not exists(biomol) and exists(i.biomol):
         biomol = i.biomol
     else:
         # construct a `Biomolecule` object from the input PDB mmCIF file
 
-        filepath = pdb_input.mmcif_filepath
-        file_id = os.path.splitext(os.path.basename(filepath))[0]
         assert os.path.exists(filepath), f"PDB input file `{filepath}` does not exist."
 
         mmcif_object = mmcif_parsing.parse_mmcif_object(
@@ -2162,14 +2163,20 @@ def pdb_input_to_molecule_input(
         assert exists(
             i.cropping_config
         ), "A cropping configuration must be provided during training."
-        biomol = biomol.crop(
-            contiguous_weight=i.cropping_config["contiguous_weight"],
-            spatial_weight=i.cropping_config["spatial_weight"],
-            spatial_interface_weight=i.cropping_config["spatial_interface_weight"],
-            n_res=i.cropping_config["n_res"],
-            chain_1=i.chains[0],
-            chain_2=i.chains[1],
-        )
+        try:
+            assert exists(i.chains), "Chain IDs must be provided for cropping during training."
+            chain_id_1, chain_id_2 = i.chains
+
+            biomol = biomol.crop(
+                contiguous_weight=i.cropping_config["contiguous_weight"],
+                spatial_weight=i.cropping_config["spatial_weight"],
+                spatial_interface_weight=i.cropping_config["spatial_interface_weight"],
+                n_res=i.cropping_config["n_res"],
+                chain_1=chain_id_1 if chain_id_1 else None,
+                chain_2=chain_id_2 if chain_id_2 else None,
+            )
+        except Exception as e:
+            raise ValueError(f"Failed to crop the biomolecule for input {file_id} due to: {e}")
 
     # retrieve features directly available within the `Biomolecule` object