#| The elements through krypton, by atomic number minus one.
ELEMENTS = ["H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na", "Mg", "Al", "Si", "P", "S", "Cl", "Ar", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe", "Co", "Ni", "Cu", "Zn", "Ga", "Ge", "As", "Se", "Br", "Kr"]
#| A molecule's atoms: symbols, atomic numbers, and positions (one row per atom).
record Geometry {
symbols: [Str]
z: Arr
xyz: Arr
}
#| The atomic number of an element symbol, in either case ("o" is oxygen).
def atomic_number(sym: Str) {
s = sym[:1].upper() + sym[1:].lower()
if s not in ELEMENTS {
stop(f"unknown element {sym!r}: the elements known here are H to Kr")
}
return index(ELEMENTS, s) + 1
}
#| The atoms of `text`: XYZ with its two header lines, or atoms alone,
#| separated by newlines or semicolons.
def parse_xyz(text: Str) -> Geometry {
lines = text.splitlines()
first = 0
declared = -1
# the XYZ header: a count, then a comment line, which is not read
if len(lines) >= 2 and lines[0].strip().isdigit() {
declared = int(lines[0])
first = 2
}
symbols = []
z = []
xyz = []
for i from first to len(lines) {
for atom in lines[i].split(";") {
f = atom.split()
if len(f) == 0 {
continue
}
if len(f) != 4 {
stop(f"line {i + 1:d}: an atom is a symbol and three numbers, not {atom.strip()!r}")
}
symbols = append(symbols, f[0])
z = append(z, atomic_number(f[0]))
xyz = append(xyz, [float(f[1]), float(f[2]), float(f[3])])
}
}
if declared >= 0 and declared != len(symbols) {
stop(f"the header says {declared:d} atoms and the text has {len(symbols):d}")
}
return Geometry { symbols: symbols, z: z, xyz: xyz }
}
test "water from an XYZ file" {
m = parse_xyz("3\nwater, from the NIST geometry\nO 0.000000 0.000000 0.117300\nH 0.000000 0.757200 -0.469200\nH 0.000000 -0.757200 -0.469200\n")
assert_true(m.symbols == ["O", "H", "H"])
assert_true(all(m.z == [8.0, 1.0, 1.0]))
assert_true(all(shape(m.xyz) == [3.0, 3.0]))
assert_true(m.xyz[0, 2] == 0.1173 and m.xyz[1, 1] == 0.7572 and m.xyz[2, 1] == -0.7572)
}
test "the same atoms as PySCF writes them" {
m = parse_xyz("o 0 0 0.1173; H 0 0.7572 -0.4692; h 0 -0.7572 -0.4692")
assert_true(m.symbols == ["o", "H", "h"])
assert_true(all(m.z == [8.0, 1.0, 1.0]))
assert_true(m.xyz[2, 2] == -0.4692)
}
test "blank lines and a carbon" {
m = parse_xyz("\nC 1.5 -2 3e-1\n\nKr 0 0 0\n")
assert_true(all(m.z == [6.0, 36.0]))
assert_true(m.xyz[0, 0] == 1.5 and m.xyz[0, 2] == 0.3)
}
# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors