from pathlib import Path
from typing import Any, List, Optional, Tuple, Union
import numpy as np
from kliff.error import InputError, KeyNotFoundError
[docs]def read_extxyz(
filename: Path,
) -> Tuple[
np.ndarray,
List[str],
np.ndarray,
List[bool],
Union[float, None],
Union[np.ndarray, None],
Union[List[float], None],
]:
"""
Read atomic configuration stored in extended xyz file_format.
Args:
filename: filename to the extended xyz file
Returns:
cell: 3x3 array, supercell lattice vectors
species: species of atoms
coords: Nx3 array, coordinates of atoms
PBC: periodic boundary conditions
energy: potential energy of the configuration; `None` if not provided in file
forces: Nx3 array, forces on atoms; `None` if not provided in file
stress: 1D array of size 6, stress on the cell in Voigt notation; `None` if not
provided in file
"""
with open(filename, "r") as fin:
lines = fin.readlines()
try:
natoms = int(lines[0].split()[0])
except ValueError as e:
raise InputError(f"{e}.\nCorrupted data at line 1 of file {filename}.")
# lattice vector
line = lines[1].replace("'", '"')
cell = _parse_key_value(line, "Lattice", "float", 9, filename)
cell = np.reshape(cell, (3, 3))
# PBC
PBC = _parse_key_value(line, "PBC", "int", 3, filename)
# energy is optional
try:
in_quotes = _check_in_quotes(line, "Energy", filename)
energy = _parse_key_value(line, "Energy", "float", 1, filename, in_quotes)[
0
]
except KeyNotFoundError:
energy = None
# stress is optional
try:
stress = _parse_key_value(line, "Stress", "float", 6, filename)
except KeyNotFoundError:
stress = None
# body, species symbol, x, y, z (and fx, fy, fz if provided)
species = []
coords = []
forces = []
# if forces provided
line = lines[2].strip().split()
if len(line) == 4:
has_forces = False
elif len(line) == 7:
has_forces = True
else:
raise InputError(f"Corrupted data at line 3 of file {filename}.")
try:
num_lines = 0
for line in lines[2:]:
num_lines += 1
line = line.strip().split()
if len(line) != 4 and len(line) != 7:
raise InputError(
f'Corrupted data at line {num_lines + 3} of file "{filename}".'
)
if has_forces:
symbol, x, y, z, fx, fy, fz = line
species.append(symbol.lower().capitalize())
coords.append([float(x), float(y), float(z)])
forces.append([float(fx), float(fy), float(fz)])
else:
symbol, x, y, z = line
species.append(symbol.lower().capitalize())
coords.append([float(x), float(y), float(z)])
except ValueError as e:
raise InputError(
f"{e}.\nCorrupted data at line {num_lines + 3} of file {filename}."
)
if num_lines != natoms:
raise InputError(
f"Corrupted data file {filename}. Number of atoms is {natoms}, "
f"whereas number of data lines is {num_lines}."
)
coords = np.asarray(coords)
if has_forces:
forces = np.asarray(forces)
else:
forces = None
return cell, species, coords, PBC, energy, forces, stress
[docs]def write_extxyz(
filename: Path,
cell: np.ndarray,
species: List[str],
coords: np.ndarray,
PBC: List[bool],
energy: Optional[float] = None,
forces: Optional[np.ndarray] = None,
stress: Optional[List[float]] = None,
):
"""
Write configuration info to a file in extended xyz file_format.
Args:
filename: filename to the extended xyz file
cell: 3x3 array, supercell lattice vectors
species: species of atoms
coords: Nx3 array, coordinates of atoms
PBC: periodic boundary conditions
energy: potential energy of the configuration; If `None`, not write to file
forces: Nx3 array, forces on atoms; If `None`, not write to file
stress: 1D array of size 6, stress on the cell in Voigt notation; If `None`,
not write to file
"""
with open(filename, "w") as fout:
# first line (number of atoms)
natoms = len(species)
fout.write("{}\n".format(natoms))
# second line
fout.write('Lattice="')
for i, line in enumerate(cell):
for j, item in enumerate(line):
if i == 2 and j == 2:
fout.write('{:.15g}" '.format(item))
else:
fout.write("{:.15g} ".format(item))
PBC = [int(i) for i in PBC]
fout.write('PBC="{} {} {}" '.format(PBC[0], PBC[1], PBC[2]))
if energy is not None:
fout.write('Energy="{:.15g}" '.format(energy))
if stress is not None:
fout.write('Stress="')
for i, s in enumerate(stress):
if i == 5:
fout.write('{:.15g}" '.format(s))
else:
fout.write("{:.15g} ".format(s))
properties = "Properties=species:S:1:pos:R:3"
if forces is not None:
properties += ":for:R:3\n"
else:
properties += "\n"
fout.write(properties)
# body
for i in range(natoms):
fout.write("{:2s} ".format(species[i]))
fout.write(
"{:23.15e} {:23.15e} {:23.15e} ".format(
coords[i][0], coords[i][1], coords[i][2]
)
)
if forces is not None:
fout.write(
"{:23.15e} {:23.15e} {:23.15e}".format(
forces[i][0], forces[i][1], forces[i][2]
)
)
fout.write("\n")
def _parse_key_value(
line: str, key: str, dtype: str, size: int, filename: Path, in_quotes: bool = True
) -> List[Any]:
"""
Given key, parse a string like ``other stuff key="value" other stuff`` to get value.
If there is not space in value, the quotes `"` can be omitted.
Args:
line: The string line.
key: Keyword to parse.
dtype: Expected data type of value, `int` or `float`.
size: Expected size of value.
filename: File name where the line comes from.
Returns:
Values associated with key.
"""
line = line.strip()
key = _check_key(line, key, filename)
try:
value = line[line.index(key) :]
if in_quotes:
value = value[value.index('"') + 1 :]
value = value[: value.index('"')]
else:
value = value[value.index("=") + 1 :]
value = value.lstrip(" ")
value += " " # add an whitespace at end in case this is the last key
value = value[: value.index(" ")]
value = value.split()
except Exception as e:
raise InputError(f"{e}.\nCorrupted {key} data at line 2 of file {filename}.")
if len(value) != size:
raise InputError(
f"Incorrect size of {key} at line 2 of file {filename};\n"
f"required: {size}, provided: {len(value)}. Possibly, the quotes not match."
)
try:
if dtype == "float":
value = [float(i) for i in value]
elif dtype == "int":
value = [int(i) for i in value]
except Exception as e:
raise InputError(f"{e}.\nCorrupted {key} data at line 2 of file {filename}.")
return value
def _check_key(line, key, filename):
"""
Check whether a key or its lowercase counter part is in line.
"""
if key not in line:
key_lower = key.lower()
if key_lower not in line:
raise KeyNotFoundError(f"{key} not found at line 2 of file {filename}.")
else:
key = key_lower
return key
def _check_in_quotes(line, key, filename):
"""
Check whether ``key=value`` or ``key="value"`` in line.
"""
key = _check_key(line, key, filename)
value = line[line.index(key) :]
value = value[value.index("=") + 1 :]
value = value.lstrip(" ")
if value[0] == '"':
return True
else:
return False