Add updated template indices for distograms and token center atoms in life.py (#84)

* Add updated template indices for distograms and token center atoms in `life.py`

* Update amino_acid_constants.py

* Update dna_constants.py

* Update ligand_constants.py

* Update rna_constants.py

* Update life.py

* Update 7a4d-assembly1.cif

* Update life.py

Author: amorehead

alphafold3_pytorch/common/amino_acid_constants.py

@@ -1,5 +1,7 @@
 """Amino acid constants used in AlphaFold."""
 
+import numpy as np
+
 from typing import Final
 
 # This mapping is used when we need to store atom data in a format that requires
@@ -173,5 +175,21 @@
     ],
     "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH", "", ""],
     "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""],
-    "UNK": ["", "", "", "", "", "", "", "", "", "", "", "", "", ""],
+    "UNK": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""],
 }
+
+restype_atom47_to_compact_atom = np.zeros([21, 47], dtype=int)
+
+
+def _make_constants():
+    """Fill the array(s) above."""
+    for restype, restype_letter in enumerate(restypes):
+        resname = restype_1to3[restype_letter]
+        for compact_atomidx, atomname in enumerate(restype_name_to_compact_atom_names[resname]):
+            if not atomname:
+                continue
+            atomtype = atom_order[atomname]
+            restype_atom47_to_compact_atom[restype, atomtype] = compact_atomidx
+
+
+_make_constants()

alphafold3_pytorch/common/dna_constants.py

@@ -1,5 +1,7 @@
 """Deoxyribonucleic acid (DNA) constants used in AlphaFold."""
 
+import numpy as np
+
 from typing import Final
 
 from alphafold3_pytorch.common import amino_acid_constants, rna_constants
@@ -209,29 +211,45 @@
         "",
     ],
     "DN": [
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
+        "OP3",
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "C1'",
+        "N9",
+        "C8",
+        "N7",
+        "C5",
+        "C6",
+        "N6",
+        "N1",
+        "C2",
+        "N3",
+        "C4",
         "",
         "",
     ],
 }
+
+restype_atom47_to_compact_atom = np.zeros([5, 47], dtype=int)
+
+
+def _make_constants():
+    """Fill the array(s) above."""
+    for restype, restype_letter in enumerate(restypes):
+        resname = restype_1to3[restype_letter]
+        for atomname in restype_name_to_compact_atom_names[resname]:
+            if not atomname:
+                continue
+            atomtype = atom_order[atomname]
+            compact_atom_idx = restype_name_to_compact_atom_names[resname].index(atomname)
+            restype_atom47_to_compact_atom[restype, atomtype] = compact_atom_idx
+
+_make_constants()

alphafold3_pytorch/common/ligand_constants.py

@@ -1,5 +1,7 @@
 """Ligand constants used in AlphaFold."""
 
+import numpy as np
+
 from typing import Final
 
 from alphafold3_pytorch.common import amino_acid_constants, dna_constants
@@ -62,11 +64,14 @@
 
 
 # All ligand residues are mapped to the unknown amino acid type index (:= 20).
-restypes = ["UNL"]
+restypes = ["X"]
 min_restype_num = len(amino_acid_constants.restypes)  # := 20.
 restype_order = {restype: min_restype_num + i for i, restype in enumerate(restypes)}
 restype_num = len(amino_acid_constants.restypes)  # := 20.
 
+
+restype_1to3 = {"X": "UNL"}
+
 BIOMOLECULE_CHAIN: Final[str] = "other"
 POLYMER_CHAIN: Final[str] = "non-polymer"
 
@@ -89,3 +94,19 @@
 restype_name_to_compact_atom_names = {
     "UNL": atom_types,
 }
+
+restype_atom47_to_compact_atom = np.zeros([1, 47], dtype=int)
+
+
+def _make_constants():
+    """Fill the array(s) above."""
+    for restype, restype_letter in enumerate(restypes):
+        resname = restype_1to3[restype_letter]
+        for atomname in restype_name_to_compact_atom_names[resname]:
+            if not atomname:
+                continue
+            atomtype = atom_order[atomname]
+            compact_atom_idx = restype_name_to_compact_atom_names[resname].index(atomname)
+            restype_atom47_to_compact_atom[restype, atomtype] = compact_atom_idx
+
+_make_constants()

alphafold3_pytorch/common/rna_constants.py

@@ -1,5 +1,7 @@
 """Ribonucleic acid (RNA) constants used in AlphaFold."""
 
+import numpy as np
+
 from typing import Final
 
 from alphafold3_pytorch.common import amino_acid_constants
@@ -201,29 +203,45 @@
         "",
     ],
     "N": [
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
-        "",
+        "OP3",
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "O2'",
+        "C1'",
+        "N9",
+        "C8",
+        "N7",
+        "C5",
+        "C6",
+        "N6",
+        "N1",
+        "C2",
+        "N3",
+        "C4",
         "",
     ],
 }
+
+restype_atom47_to_compact_atom = np.zeros([5, 47], dtype=int)
+
+
+def _make_constants():
+    """Fill the array(s) above."""
+    for restype, restype_letter in enumerate(restypes):
+        resname = restype_1to3[restype_letter]
+        for atomname in restype_name_to_compact_atom_names[resname]:
+            if not atomname:
+                continue
+            atomtype = atom_order[atomname]
+            compact_atom_idx = restype_name_to_compact_atom_names[resname].index(atomname)
+            restype_atom47_to_compact_atom[restype, atomtype] = compact_atom_idx
+
+_make_constants()