"""Utility functions for the handling of the input files"""
from collections.abc import Iterable
from typing import Optional, Union
from aiida import orm
import numpy as np
[docs]def generate_lammps_structure(
structure: orm.StructureData,
atom_style: str = "atomic",
charge_dict: Optional[dict[str, float]] = None,
round_dp: Optional[float] = None,
docstring: str = "generated by aiida_lammps",
) -> Union[str, np.array]:
"""Creation of the structure file content.
As allowing the users to create their lattices using LAMMPS' would be too
complex, one must ensure that the aiida StructureData is written to file in
a format that is compatible to LAMMPS.
In the case of non-orthogonal structures, this will take care of generating
a triclinic cell compatible with what LAMMPS expects.
:param structure: the structure to use in the simulation
:type structure: orm.StructureData
:param atom_style: treatment of the particles according to lammps, defaults to 'atomic'
:type atom_style: str, optional
:param charge_dict: dictionary with the charge for the particles, defaults to None
:type charge_dict: dict, optional
:param round_dp: precision to which to round the positions, defaults to None
:type round_dp: float, optional
:param docstring: header for the structure file, defaults to 'generated by aiida_lammps'
:type docstring: str, optional
:raises ValueError: if the atom_style does not belong to either 'atomic' or 'charge'
:return: the structure file content, the transformation matrix applied to
the structure cell and coordinates
:rtype: Union[str, np.array]
"""
# pylint: disable=too-many-locals
if atom_style not in ["atomic", "charge"]:
raise ValueError(
f"atom_style must be in ['atomic', 'charge'], not '{atom_style}'"
)
if charge_dict is None:
charge_dict = {}
# mapping of atom kind_name to id number
kind_name_id_map = {}
for site in structure.sites:
if site.kind_name not in kind_name_id_map:
kind_name_id_map[site.kind_name] = len(kind_name_id_map) + 1
# mapping of atom kind_name to mass
kind_mass_dict = {kind.name: kind.mass for kind in structure.kinds}
filestring = ""
filestring += f"# {docstring}\n\n"
filestring += f"{len(structure.sites)} atoms\n"
filestring += f"{len(kind_name_id_map)} atom types\n\n"
cell, coord_transform = _transform_cell(structure.cell)
positions = np.transpose(
np.dot(
coord_transform, np.transpose([_site.position for _site in structure.sites])
)
)
if round_dp:
cell = np.round(cell, round_dp) + 0.0
positions = np.round(positions, round_dp) + 0.0
filestring += f"0.0 {cell[0, 0]:20.10f} xlo xhi\n"
filestring += f"0.0 {cell[1, 1]:20.10f} ylo yhi\n"
filestring += f"0.0 {cell[2, 2]:20.10f} zlo zhi\n"
filestring += (
f"{cell[1, 0]:20.10f} {cell[2, 0]:20.10f} {cell[2, 1]:20.10f} xy xz yz\n\n"
)
filestring += "Masses\n\n"
for kind_name in sorted(kind_name_id_map.keys()):
filestring += (
f"{kind_name_id_map[kind_name]} {kind_mass_dict[kind_name]:20.10f} \n"
)
filestring += "\n"
filestring += "Atoms\n\n"
for site_index, (pos, site) in enumerate(zip(positions, structure.sites)):
kind_id = kind_name_id_map[site.kind_name]
if atom_style == "atomic":
filestring += f"{site_index + 1} {kind_id}"
filestring += f" {pos[0]:20.10f} {pos[1]:20.10f} {pos[2]:20.10f}\n"
if atom_style == "charge":
charge = charge_dict.get(site.kind_name, 0.0)
filestring += f"{site_index + 1} {kind_id} {charge}"
filestring += f" {pos[0]:20.10f} {pos[1]:20.10f} {pos[2]:20.10f}\n"
return filestring, coord_transform
[docs]def flatten(full_list: list) -> list:
"""Flattens a list of list into a flat list.
:param full_list: list of lists to be flattened
:type full_list: list
:yield: flattened list
:rtype: list
"""
for element in full_list:
if isinstance(element, Iterable) and not isinstance(element, (str, bytes)):
yield from flatten(element)
else:
yield element
[docs]def convert_to_str(value) -> str:
"""convert True/False to yes/no and all values to strings"""
if isinstance(value, bool):
if value:
return "yes"
return "no"
return str(value)
[docs]def _convert_values(value) -> str:
if isinstance(value, (tuple, list)):
return " ".join([convert_to_str(v) for v in value])
return convert_to_str(value)
[docs]def join_keywords(dct, ignore=None) -> str:
"""join a dict of {keyword: value, ...} into a string 'keyword value ...'
value can be a single value or a list/tuple of values
"""
ignore = ignore if ignore else []
return " ".join(
[
f"{k} {_convert_values(dct[k])}"
for k in sorted(dct.keys())
if k not in ignore
]
)
