HOTFIX: add atom_types to CIF, contributed by @ggrigoryanv

chroma/data/system.py

@@ -2650,6 +2650,27 @@ def to_CIF_string(self):
         return cif_str
 
     def _write_cif(self, f):
+        # fmt: off
+        _specials_atom_names = [
+            "MG", "CL", "FE", "ZN", "MN", "NI", "SE", "CU", "BR", "CO", "AS",
+            "BE", "RU", "RB", "ZR", "OS", "SR", "GD", "MO", "AU", "AG", "PT",
+            "AL", "XE", "BE", "CS", "EU", "IR", "AM", "TE", "BA", "SB"
+        ]
+        # fmt: on
+        _ambiguous_atom_names = ["CA", "CD", "NA", "HG", "PB"]
+
+        def _guess_type(atom_name, res_name):
+            if len(atom_name) > 0 and atom_name[0] == '"':
+                atom_name = atom_name.replace('"', "")
+            if atom_name[:2] in _specials_atom_names:
+                return atom_name[:2]
+            else:
+                if atom_name in _ambiguous_atom_names and res_name == atom_name:
+                    return atom_name
+                elif atom_name == "UNK":
+                    return "X"
+            return atom_name[:1]
+
         entry_id = self.name.strip()
         if entry_id == "":
             entry_id = "system"
@@ -2772,6 +2793,7 @@ def _write_cif(self, f):
                 "pdbx_PDB_model_num",
                 "auth_seq_id",
                 "auth_asym_id",
+                "type_symbol",
             ],
         )
         idx = -1
@@ -2809,7 +2831,9 @@ def _write_cif(self, f):
                                 f"{sp.atom_site_token(residue.icode)} {coor[0]:g} {coor[1]:g} {coor[2]:g} "
                             )
                             f.write(f"{coor[3]:g} {coor[4]:g} {model_index} ")
-                            f.write(f"{authresid} {authchainid}\n")
+                            f.write(
+                                f"{authresid} {authchainid} {_guess_type(atom.name, residue.name)}\n"
+                            )
             self.swap_model(model_index)
         f.write("#\n")