add: support for atomic number instead of element symbols

pysisyphus/Geometry.py

@@ -23,10 +23,11 @@
 from pysisyphus.constants import BOHR2ANG
 from pysisyphus.hessian_proj import get_hessian_projector, inertia_tensor
 from pysisyphus.elem_data import (
-    MASS_DICT,
-    ISOTOPE_DICT,
     ATOMIC_NUMBERS,
     COVALENT_RADII as CR,
+    INV_ATOMIC_NUMBERS,
+    ISOTOPE_DICT,
+    MASS_DICT,
     VDW_RADII as VDWR,
 )
 from pysisyphus.helpers_pure import (
@@ -57,6 +58,22 @@
 from pysisyphus.xyzloader import make_xyz_str
 
 
+def normalize_atoms(atoms) -> tuple[str]:
+    atomic_numbers = set(INV_ATOMIC_NUMBERS.keys())
+    _atoms = list()
+    for atom in atoms:
+        try:
+            atom_int = int(atom)
+            if atom_int in atomic_numbers:
+                atom = INV_ATOMIC_NUMBERS[atom_int]
+        except ValueError:
+            pass
+        # Was atom.capitalize() before ...
+        atom = atom.lower()
+        _atoms.append(atom)
+    return tuple(_atoms)
+
+
 class Geometry:
     coord_types = {
         "cart": None,
@@ -121,7 +138,7 @@ def __init__(
         name : str, optional
             Verbose name of the geometry, e.g. methanal or water. Used for printing
         """
-        self.atoms = tuple([atom.capitalize() for atom in atoms])
+        self.atoms = normalize_atoms(atoms)
         # self._coords always holds cartesian coordinates.
         self._coords = np.array(coords, dtype=float).flatten()
         assert self._coords.size == (3 * len(self.atoms)), (
@@ -208,8 +225,10 @@ def moving_atoms_jmol(self):
 
     @property
     def sum_formula(self):
-        unique_atoms = sorted(set(self.atoms))
-        counter = Counter(self.atoms)
+        atoms = self.atoms
+        atoms = [atom.capitalize() for atom in atoms]
+        unique_atoms = sorted(set(atoms))
+        counter = Counter(atoms)
         atoms = list()
         num_strs = list()
 

pysisyphus/io/pdb.py

@@ -212,7 +212,7 @@ def atoms_coords_to_pdb_str(atoms, coords, fragments=None, resname="", conect=Tr
     serial = 1
     for resSeq, fragment in enumerate(fragments, 1):
         for id_ in fragment:
-            name = atoms[id_]
+            name = atoms[id_].capitalize()
             xyz = coords3d_ang[id_]
             line = hetatm_fmt.format(
                 serial,

tests/test_replacements/test_replacements.py

@@ -30,4 +30,4 @@ def test_replace_atom_with_atom():
     geom = geom_loader("lib:methane.xyz")
     repl_geom = Geometry(("I",), (0.0, 0.0, 0.0))
     union = replace_atom(geom, 3, repl_geom, 0)
-    assert union.atoms[3] == "I"
+    assert union.atoms[3] == "i"