"""Utilities for molecular geometries.
Notation:
- ``N``: number of atoms.
"""
from __future__ import annotations
import numpy as np
from numpy.typing import ArrayLike, NDArray
atomic_weights = {
"H": 1.008,
"He": 4.002602,
"Li": 6.94,
"Be": 9.0121831,
"B": 10.81,
"C": 12.011,
"N": 14.007,
"O": 15.999,
"F": 18.998403163,
"Ne": 20.1797,
"Na": 22.98976928,
"Mg": 24.305,
"Al": 26.9815385,
"Si": 28.085,
"P": 30.973761998,
"S": 32.06,
"Cl": 35.45,
"Ar": 39.948,
"K": 39.0983,
"Ca": 40.078,
"Sc": 44.955908,
"Ti": 47.867,
"V": 50.9415,
"Cr": 51.9961,
"Mn": 54.938044,
"Fe": 55.845,
"Co": 58.933194,
"Ni": 58.6934,
"Cu": 63.546,
"Zn": 65.38,
"Ga": 69.723,
"Ge": 72.63,
"As": 74.921595,
"Se": 78.971,
"Br": 79.904,
"Kr": 83.798,
"Rb": 85.4678,
"Sr": 87.62,
"Y": 88.90584,
"Zr": 91.224,
"Nb": 92.90637,
"Mo": 95.95,
"Tc": 97.90721,
"Ru": 101.07,
"Rh": 102.9055,
"Pd": 106.42,
"Ag": 107.8682,
"Cd": 112.414,
"In": 114.818,
"Sn": 118.71,
"Sb": 121.76,
"Te": 127.6,
"I": 126.90447,
"Xe": 131.293,
"Cs": 132.90545196,
"Ba": 137.327,
"La": 138.90547,
"Ce": 140.116,
"Pr": 140.90766,
"Nd": 144.242,
"Pm": 144.91276,
"Sm": 150.36,
"Eu": 151.964,
"Gd": 157.25,
"Tb": 158.92535,
"Dy": 162.5,
"Ho": 164.93033,
"Er": 167.259,
"Tm": 168.93422,
"Yb": 173.045,
"Lu": 174.9668,
"Hf": 178.49,
"Ta": 180.94788,
"W": 183.84,
"Re": 186.207,
"Os": 190.23,
"Ir": 192.217,
"Pt": 195.084,
"Au": 196.966569,
"Hg": 200.592,
"Tl": 204.38,
"Pb": 207.2,
"Bi": 208.9804,
"Po": 209.0,
"At": 210.0,
"Rn": 222.0,
"Fr": 223.0,
"Ra": 226.0,
"Ac": 227.0,
"Th": 232.0377,
"Pa": 231.03588,
"U": 238.02891,
"Np": 237.0,
"Pu": 244.0,
"Am": 243.0,
"Cm": 247.0,
"Bk": 247.0,
"Cf": 251.0,
"Es": 252.0,
"Fm": 257.0,
"Md": 258.0,
"No": 259.0,
"Lr": 262.0,
"Rf": 267.0,
"Db": 268.0,
"Sg": 271.0,
"Bh": 274.0,
"Hs": 269.0,
"Mt": 276.0,
"Ds": 281.0,
"Rg": 281.0,
"Cn": 285.0,
"Nh": 286.0,
"Fl": 289.0,
"Mc": 288.0,
"Lv": 293.0,
"Ts": 294.0,
"Og": 294.0,
}
[docs]
def get_geom(xyzstr: str) -> tuple[list[str], NDArray[np.float64], NDArray[np.float64]]:
"""Get the chemical symbols, positions and masses from an XYZ string.
Only the lines of the form ``symbol x y z`` whose first token is a chemical
symbol in ``atomic_weights`` are read. Other lines (e.g. the atom count and
the comment line of the XYZ format) are skipped.
Args:
xyzstr: Contents of an XYZ file.
Returns:
``(symbols, posi, mass)``, where
- ``symbols``: ``(N,)`` chemical symbols.
- ``posi``: ``(N, 3)`` positions, in the same units as ``xyzstr``.
- ``mass``: ``(N,)`` atomic masses [amu].
"""
mass = []
posi = []
symbols = []
for line in xyzstr.splitlines():
line = line.strip()
if line == "":
continue
line_split = line.split()
if line_split[0] in atomic_weights.keys():
symbol, x, y, z = line.split()
symbols.append(symbol)
mass.append(atomic_weights[symbol])
posi.append([float(s) for s in [x, y, z]])
posi = np.array(posi)
mass = np.array(mass)
return symbols, posi, mass
[docs]
def get_mass_center(posi: ArrayLike, mass: ArrayLike) -> NDArray[np.float64]:
"""Get the center of mass.
Args:
posi: ``(N, 3)`` positions.
mass: ``(N,)`` atomic masses.
Returns:
``(3,)`` center of mass.
"""
return np.sum(posi * mass[:, None], axis=0) / np.sum(mass)
[docs]
def get_inertia(posi: ArrayLike, mass: ArrayLike) -> NDArray[np.float64]:
r"""Get the inertia tensor about the origin.
.. math::
I = \sum_a m_a ((r_a^T r_a) I_3 - r_a r_a^T)
Args:
posi: ``(N, 3)`` positions.
mass: ``(N,)`` atomic masses.
Returns:
``(3, 3)`` inertia tensor.
Note:
The tensor is computed about the origin. To obtain the inertia tensor
about the center of mass, subtract :func:`get_mass_center` from
``posi`` in advance.
"""
n_atom = len(mass)
xx = np.einsum("ax,ax->a", posi, posi)
xy = np.einsum("ax,ay->axy", posi, posi)
inertia = np.zeros((n_atom, 3, 3))
diag = [0, 1, 2]
inertia[:, diag, diag] = xx[:, np.newaxis]
inertia -= xy
inertia = np.einsum("a,axy->xy", mass, inertia)
return inertia