Address non-standard atoms parsing for polymer residues (#212)

* Update README.md

* Update amino_acid_constants.py

* Update dna_constants.py

* Update rna_constants.py

Author: amorehead

README.md

@@ -52,6 +52,10 @@ A visualization of the molecules of life used in the repository can be seen and
 
 - <a href="https://github.com/dhuvik">Dhuvi</a> for fixing a bug related to metal ion molecule ID assignment for `Alphafold3Inputs`!
 
+- Tom (from the Discord channel) for identifying a discrepancy between this codebase's distogram and template unit vector computations and those of OpenFold (and <a href="https://github.com/vandrw">Andrei</a> for addressing these issues)!
+
+- <a href="https://github.com/CyaNine">Kaihui</a> for identifying a bug in how non-standard atoms were handled in polymer residues!
+
 - <a href="https://github.com/patrick-kidger">Patrick</a> for <a href="https://docs.kidger.site/jaxtyping/">jaxtyping</a>, <a href="https://github.com/fferflo">Florian</a> for <a href="https://github.com/fferflo/einx">einx</a>, and of course, <a href="https://github.com/arogozhnikov">Alex</a> for <a href="https://einops.rocks/">einops</a>
 
 ## Install

alphafold3_pytorch/common/amino_acid_constants.py

@@ -45,7 +45,7 @@
     "CZ3",
     "NZ",
     # "OXT",  # NOTE: This often appears in mmCIF files, but it will not be used for any amino acid type in AlphaFold.
-    "_",
+    "ATM",  # NOTE: This represents a catch-all atom type for non-standard or modified residues.
     "_",
     "_",
     "_",
@@ -57,7 +57,7 @@
     "_",
     "_",  # 10 null types.
 ]
-element_types = [atom_type[0] for atom_type in atom_types]
+element_types = [atom_type if atom_type == "ATM" else atom_type[0] for atom_type in atom_types]
 atom_types_set = set(atom_types)
 atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
 atom_type_num = len(atom_types)  # := 37 + 10 null types := 47.

alphafold3_pytorch/common/dna_constants.py

@@ -39,7 +39,7 @@
     "N2",
     "O4",
     "C7",
-    "_",
+    "ATM",  # NOTE: This represents a catch-all atom type for non-standard or modified residues.
     "_",
     "_",
     "_",
@@ -59,7 +59,7 @@
     "_",
     "_",  # 19 null types.
 ]
-element_types = [atom_type[0] for atom_type in atom_types]
+element_types = [atom_type if atom_type == "ATM" else atom_type[0] for atom_type in atom_types]
 atom_types_set = set(atom_types)
 atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
 atom_type_num = len(atom_types)  # := 28 + 19 null types := 47.

alphafold3_pytorch/common/rna_constants.py

@@ -39,7 +39,7 @@
     "O6",
     "N2",
     "O4",
-    "_",
+    "ATM",  # NOTE: This represents a catch-all atom type for non-standard or modified residues.
     "_",
     "_",
     "_",
@@ -59,7 +59,7 @@
     "_",
     "_",  # 19 null types.
 ]
-element_types = [atom_type[0] for atom_type in atom_types]
+element_types = [atom_type if atom_type == "ATM" else atom_type[0] for atom_type in atom_types]
 atom_types_set = set(atom_types)
 atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
 atom_type_num = len(atom_types)  # := 28 + 19 null types := 47.