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() diff --git a/pysisyphus/io/pdb.py b/pysisyphus/io/pdb.py index 02d7709fb..20b8ce3b4 100644 --- a/pysisyphus/io/pdb.py +++ b/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, diff --git a/tests/test_replacements/test_replacements.py b/tests/test_replacements/test_replacements.py index be8c84266..5dccf061d 100644 --- a/tests/test_replacements/test_replacements.py +++ b/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"