From e84c142528d1ebb6265729ad82c48bb1c39bc059 Mon Sep 17 00:00:00 2001
From: Johannes Steinmetzer <johannes.steinmetzer@uni-jena.de>
Date: Thu, 5 Oct 2023 09:57:15 +0200
Subject: [PATCH] add: support for atomic number instead of element symbols

---
 pysisyphus/Geometry.py | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/pysisyphus/Geometry.py b/pysisyphus/Geometry.py
index 1d6b2b052..2e46efd67 100644
--- a/pysisyphus/Geometry.py
+++ b/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()