Handle edge case for unknown polymer residues (#228)

* Update cli.py

* Update amino_acid_constants.py

* Update dna_constants.py

* Update rna_constants.py

* Update inputs.py

Author: amorehead

alphafold3_pytorch/cli.py

@@ -16,15 +16,15 @@
 @click.command()
 @click.option('-ckpt', '--checkpoint', type = str, help = 'path to alphafold3 checkpoint')
 @click.option('-p', '--protein', type = str, help = 'one protein sequence')
-@click.option('-o', '--output', type = str, help = 'output path', default = 'output.mmcif')
+@click.option('-o', '--output', type = str, help = 'output path', default = 'output.cif')
 def cli(
     checkpoint: str,
     protein: str,
     output: str
 ):
 
     checkpoint_path = Path(checkpoint)
-    assert checkpoint_path.exists(), f'alphafold3 checkpoint must exist at {str(checkpoint_path)}'
+    assert checkpoint_path.exists(), f'AlphaFold 3 checkpoint must exist at {str(checkpoint_path)}'
 
     alphafold3_input = Alphafold3Input(
         proteins = [protein],
@@ -44,4 +44,4 @@ def cli(
     pdb_writer.set_structure(structure)
     pdb_writer.save(str(output_path))
 
-    print(f'mmcif saved to {str(output_path)}')
+    print(f'mmCIF file saved to {str(output_path)}')

alphafold3_pytorch/common/amino_acid_constants.py

@@ -216,7 +216,7 @@
 
 def _make_constants():
     """Fill the array(s) above."""
-    for restype, restype_letter in enumerate(restypes):
+    for restype, restype_letter in enumerate(restype_1to3.keys()):
         resname = restype_1to3[restype_letter]
         for compact_atomidx, atomname in enumerate(restype_name_to_compact_atom_names[resname]):
             if not atomname:

alphafold3_pytorch/common/dna_constants.py

@@ -255,7 +255,7 @@
 
 def _make_constants():
     """Fill the array(s) above."""
-    for restype, restype_letter in enumerate(restypes):
+    for restype, restype_letter in enumerate(restype_1to3.keys()):
         resname = restype_1to3[restype_letter]
         for atomname in restype_name_to_compact_atom_names[resname]:
             if not atomname:

alphafold3_pytorch/common/rna_constants.py

@@ -246,7 +246,7 @@
 
 def _make_constants():
     """Fill the array(s) above."""
-    for restype, restype_letter in enumerate(restypes):
+    for restype, restype_letter in enumerate(restype_1to3.keys()):
         resname = restype_1to3[restype_letter]
         for atomname in restype_name_to_compact_atom_names[resname]:
             if not atomname:

alphafold3_pytorch/inputs.py

@@ -2281,10 +2281,17 @@ def extract_canonical_molecules_from_biomolecule_chains(
                 contiguous_res_atom_mapping = np.vectorize(contiguous_res_atom_mapping.get)(
                     res_atom_mapping
                 )
-
                 res_atom_positions = atom_positions[res_index][res_atom_mask][
                     contiguous_res_atom_mapping
                 ]
+
+                num_atom_positions = len(res_atom_positions) + len(missing_atom_indices)
+                if num_atom_positions != mol.GetNumAtoms():
+                    raise ValueError(
+                        f"The number of (missing and present) atom positions ({num_atom_positions}) for residue {res} does not match the number of atoms in the RDKit molecule ({mol.GetNumAtoms()}). "
+                        "Please ensure that these input features are correctly paired. Skipping this example."
+                    )
+                
                 mol = add_atom_positions_to_mol(
                     mol,
                     res_atom_positions.reshape(-1, 3),