Source code for autojob.utils.vasp

"""VASP utilities."""

from collections.abc import Iterable
import logging
from pathlib import Path
from typing import Literal
from typing import TypeVar

logger = logging.getLogger(__name__)

_T = TypeVar("_T")




def reorder_vasp_sequence(
    to_reorder: Iterable[_T],
    dir_name: str | Path,
    *,
    direction: Literal["to_ase", "to_vasp"] = "to_ase",
) -> list[_T]:
    """Reorder a sequence to/from a VASP ordering from/to an ASE ordering."""
    logger.debug(
        f"Reordering a sequence: {type(to_reorder)} of type: {type(to_reorder[0])}"
    )
    sort_file = Path(dir_name, "ase-sort.dat")

    with Path(sort_file).open(mode="r", encoding="utf-8") as file:
        lines = file.readlines()

    # First column: the ith atom in the VASP structure corresponds to
    # the atom in the ase.Atoms object with index given by the integer in row i
    # Second column: the ith atom in the original ase.Atoms object corresponds
    # to the atom in the VASP structure with index given by the integer in row i
    column = 0 if direction == "to_vasp" else 1
    ordering = [int(line.split()[column]) for line in lines]
    reordered = [to_reorder[i] for i in ordering]

    logger.debug(
        "Successfully reordered sequence: "
        f"{[i for i, _ in enumerate(to_reorder)]!r} -> {ordering!r}"
    )
    return reordered