"""POSCAR/INCAR/KPOINTS/POTCAR preparation and provenance helpers.
This is the input-side of the dependency-free VASP interface: reading and
rewriting the four VASP input files, assembling a POTCAR with recorded
provenance, and the :func:`prepare_vasp_inputs` entry point that ties them
together. Historical authorship is documented in ``v1_runtime/NOTICE``.
"""
import bz2
import gzip
import hashlib
import logging
import lzma
import math
import os
import random
import re
import shutil
import subprocess
import tempfile
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from .._util import sha256_file, utc_now, write_json_atomic
_LOGGER = logging.getLogger(__name__)
#: The k-point centering every entry point starts from. Monkhorst-Pack is the
#: starting point rather than Gamma because the reviewed remedy ladder promotes
#: ``("centering", "Gamma")`` as the *fix* for the ``kpoints_class`` and
#: ``kpoint_shifts`` failure classes: a workflow that already starts at Gamma has
#: no such remedy left to apply.
DEFAULT_KPOINT_CENTERING = "Monkhorst-Pack"
KPOINT_CENTERINGS = ("Gamma", "Monkhorst-Pack")
[docs]
def suggested_magnetic_moments(path: str | os.PathLike[str] = "POSCAR") -> str:
"""Return the explicit comment override or a five-per-atom default.
The five Bohr magnetons per atom are a heuristic starting guess, not a
physical prediction: a deliberately high initial moment lets a spin-polarized
relaxation fall into a low-spin solution, whereas starting too low tends to keep
it there, so overestimating is the safer direction. The value of five is the
empirical default carried over from *httk* v1. Encode a per-structure choice in
the POSCAR comment as ``[MAGMOM=...]``.
"""
header = read_poscar_header(path)
override = re.search(r"\[MAGMOM=([^\]]+)\]", header.comment)
if override is not None:
return override.group(1)
return " ".join(f"{count}*5" for count in header.counts)
def _cross(left: Sequence[float], right: Sequence[float]) -> tuple[float, float, float]:
return (
left[1] * right[2] - left[2] * right[1],
left[2] * right[0] - left[0] * right[2],
left[0] * right[1] - left[1] * right[0],
)
def _dot(left: Sequence[float], right: Sequence[float]) -> float:
return sum(a * b for a, b in zip(left, right, strict=True))
[docs]
def automatic_kpoint_grid(
density: float,
*,
poscar: str | os.PathLike[str] = "POSCAR",
minimum: int = 3,
equal: bool = False,
bump: int = 0,
) -> tuple[int, int, int]:
"""Calculate a reciprocal-length automatic grid."""
if density <= 0 or minimum < 1 or bump < 0:
raise ValueError("density must be positive; minimum must be positive; bump cannot be negative")
header = read_poscar_header(poscar)
first, second, third = header.lattice
determinant = _dot(first, _cross(second, third))
if abs(determinant) < 1e-12:
raise ValueError("POSCAR lattice is singular")
scale = header.scale
if scale < 0:
scale = (-scale / abs(determinant)) ** (1.0 / 3.0)
if scale == 0:
raise ValueError("POSCAR scale cannot be zero")
reciprocal = (
_cross(second, third),
_cross(third, first),
_cross(first, second),
)
lengths = tuple(math.sqrt(_dot(row, row)) / abs(determinant * scale) for row in reciprocal)
values = tuple(max(minimum, math.ceil(density * length + 0.5) + bump) for length in lengths)
grid = (values[0], values[1], values[2])
if equal:
largest = max(grid)
return largest, largest, largest
return grid
[docs]
def write_automatic_kpoints(
grid: Sequence[int],
path: str | os.PathLike[str] = "KPOINTS",
*,
centering: str = DEFAULT_KPOINT_CENTERING,
) -> Path:
"""Write a standard automatic KPOINTS file.
The centering defaults to :data:`DEFAULT_KPOINT_CENTERING` here, in
:class:`~httk.workflow.vasp.inputs.VaspPreparationOptions`, and in the Bash bridge, so a workflow that
hits a k-point failure class still has the ``Gamma`` remedy available.
"""
values = tuple(grid)
if len(values) != 3 or any(isinstance(value, bool) or value < 1 for value in values):
raise ValueError("grid must contain three positive integers")
if centering not in KPOINT_CENTERINGS:
raise ValueError("centering must be 'Gamma' or 'Monkhorst-Pack'")
destination = Path(path)
destination.write_text(
f"Automatic mesh generated by httk\n0\n{centering}\n{values[0]} {values[1]} {values[2]}\n0 0 0\n",
encoding="utf-8",
)
return destination
[docs]
def read_incar(path: str | os.PathLike[str] = "INCAR") -> dict[str, str]:
"""Read the last value of each simple INCAR assignment."""
result: dict[str, str] = {}
for line in Path(path).read_text(encoding="utf-8").splitlines():
content = re.split(r"[#!]", line, maxsplit=1)[0]
for statement in content.split(";"):
name, separator, value = statement.partition("=")
key = name.strip().upper()
if separator and key:
result[key] = value.strip()
return result
def _incar_statement_tag(statement: str) -> str | None:
"""Return the tag one ``NAME = VALUE`` statement assigns, if it assigns one."""
name, separator, _ = statement.partition("=")
key = name.strip().upper()
return key if separator and key else None
[docs]
def update_incar(
values: Mapping[str, object],
path: str | os.PathLike[str] = "INCAR",
) -> Path:
"""Atomically replace selected INCAR tags while preserving other lines.
VASP allows several ``;``-separated assignments per line, and
:func:`read_incar` reads them all, so an update has to rewrite the individual
statements of a line rather than drop or keep the whole line: updating
``ISYM`` in ``ISPIN = 2 ; ISYM = 2`` leaves ``ISPIN = 2`` and appends the new
``ISYM``, instead of leaving a line that still assigns the old value.
"""
destination = Path(path)
normalized = {str(name).strip().upper(): str(value) for name, value in values.items()}
if not normalized or any(not name or re.fullmatch(r"[A-Z][A-Z0-9_]*", name) is None for name in normalized):
raise ValueError("INCAR updates require valid nonempty tag names")
kept: list[str] = []
for line in destination.read_text(encoding="utf-8").splitlines():
comment_match = re.search(r"[#!]", line)
content = line if comment_match is None else line[: comment_match.start()]
comment = "" if comment_match is None else line[comment_match.start() :]
statements = content.split(";")
if not any(_incar_statement_tag(item) in normalized for item in statements):
# Nothing on this line is replaced, so it survives byte for byte.
kept.append(line)
continue
surviving = [item.strip() for item in statements if _incar_statement_tag(item) not in normalized]
remainder = " ; ".join(item for item in surviving if item)
if remainder and comment:
kept.append(f"{remainder} {comment}")
elif remainder or comment:
kept.append(remainder or comment)
# Appended in tag order: the same set of updates then always produces the same
# file, whatever order the caller's mapping happened to have.
kept.extend(f"{name} = {value}" for name, value in sorted(normalized.items()))
descriptor, temporary_name = tempfile.mkstemp(prefix=f".{destination.name}.", dir=destination.parent)
temporary = Path(temporary_name)
try:
with os.fdopen(descriptor, "w", encoding="utf-8") as stream:
stream.write("\n".join(kept) + "\n")
os.replace(temporary, destination)
finally:
temporary.unlink(missing_ok=True)
return destination
def _read_potential(path: Path) -> bytes:
if path.name == "POTCAR":
return path.read_bytes()
if path.name == "POTCAR.gz":
return gzip.decompress(path.read_bytes())
if path.name == "POTCAR.bz2":
return bz2.decompress(path.read_bytes())
if path.name in {"POTCAR.xz", "POTCAR.lzma"}:
return lzma.decompress(path.read_bytes())
if path.name == "POTCAR.Z":
executable = shutil.which("gzip") or shutil.which("uncompress")
if executable is None:
raise RuntimeError("reading legacy POTCAR.Z requires gzip or uncompress")
completed = subprocess.run(
[executable, "-cd", str(path)],
check=False,
capture_output=True,
)
if completed.returncode:
raise ValueError(f"cannot decompress {path}: {completed.stderr.decode(errors='replace').strip()}")
return completed.stdout
raise ValueError(f"unsupported pseudopotential compression: {path}")
[docs]
@dataclass(frozen=True)
class PotcarChoice:
"""Which pseudopotential of a library one species was given.
The suffix policy of :func:`assemble_potcar` silently prefers one PAW variant
over another, and the choice changes the numbers a calculation produces, so
every choice is recorded: the variant directory, the suffix that selected it,
the full source path, its digest, and the ``TITEL`` the potential names
itself with.
"""
species: str
variant: str
suffix: str
source: Path
sha256: str
titel: str | None
[docs]
def as_mapping(self) -> dict[str, object]:
return {
"species": self.species,
"variant": self.variant,
"suffix": self.suffix,
"source": str(self.source),
"sha256": self.sha256,
"titel": self.titel,
}
[docs]
@dataclass(frozen=True)
class PotcarAssembly:
"""One assembled POTCAR and the provenance of every piece in it."""
path: Path
library: Path
choices: tuple[PotcarChoice, ...]
provenance: Path
[docs]
def as_mapping(self) -> dict[str, object]:
return {
"format": "httk-vasp-potcar-provenance",
"format_version": 1,
"potcar": str(self.path),
"library": str(self.library),
"assembled_at": utc_now(),
"potentials": [item.as_mapping() for item in self.choices],
}
[docs]
def assemble_potcar(
library: str | os.PathLike[str],
*,
poscar: str | os.PathLike[str] = "POSCAR",
output: str | os.PathLike[str] = "POTCAR",
suffix_preference: Iterable[str] = ("_3", "_2", "_d", "_sv", "_pv", "", "_h", "_s"),
provenance: str | os.PathLike[str] | None = None,
) -> PotcarAssembly:
"""Assemble POTCAR from explicit species and a configurable suffix policy.
The suffix policy decides which PAW variant each species gets, so the result
describes what it chose: the returned :class:`PotcarAssembly` names every
potential, and the same record is written next to the POTCAR as
``<POTCAR>.provenance.json`` unless *provenance* names another file.
"""
root = Path(library).expanduser().resolve()
header = read_poscar_header(poscar)
suffixes = tuple(suffix_preference)
pieces: list[bytes] = []
choices: list[PotcarChoice] = []
for species in header.species:
found: Path | None = None
variant = ""
for suffix in suffixes:
directory = root / f"{species}{suffix}"
for filename in ("POTCAR", "POTCAR.gz", "POTCAR.bz2", "POTCAR.xz", "POTCAR.lzma", "POTCAR.Z"):
candidate = directory / filename
if candidate.is_file():
found = candidate
variant = suffix
break
if found is not None:
break
if found is None:
raise FileNotFoundError(f"no pseudopotential found for species {species!r} below {root}")
content = _read_potential(found)
pieces.append(content)
titel = re.search(r"^\s*TITEL\s*=\s*(.+)$", content.decode("utf-8", errors="replace"), re.MULTILINE)
choices.append(
PotcarChoice(
species,
f"{species}{variant}",
variant,
found,
sha256_file(found),
None if titel is None else titel.group(1).strip(),
)
)
destination = Path(output)
destination.write_bytes(b"".join(pieces))
record = (
Path(provenance) if provenance is not None else destination.with_name(f"{destination.name}.provenance.json")
)
assembly = PotcarAssembly(destination, root, tuple(choices), record)
write_json_atomic(record, assembly.as_mapping())
_LOGGER.info(
"assembled %s from %s: %s",
destination,
root,
", ".join(f"{item.species}->{item.variant}" for item in choices),
)
return assembly
[docs]
def contcar_to_poscar(
contcar: str | os.PathLike[str] = "CONTCAR",
*,
reference: str | os.PathLike[str] = "POSCAR",
output: str | os.PathLike[str] = "POSCAR",
) -> Path:
"""Replace CONTCAR's comment with the reference POSCAR comment."""
reference_lines = Path(reference).read_text(encoding="utf-8").splitlines(keepends=True)
lines = Path(contcar).read_text(encoding="utf-8").splitlines(keepends=True)
if not reference_lines or not lines:
raise ValueError("CONTCAR and reference POSCAR must both be nonempty")
destination = Path(output)
destination.write_text(reference_lines[0].rstrip("\r\n") + "\n" + "".join(lines[1:]), encoding="utf-8")
return destination
[docs]
def normalize_poscar_handedness(path: str | os.PathLike[str] = "POSCAR") -> Path:
"""Make a left-handed POSCAR lattice right-handed without moving sites."""
destination = Path(path)
lines = destination.read_text(encoding="utf-8").splitlines()
header = read_poscar_header(destination)
determinant = _dot(header.lattice[0], _cross(header.lattice[1], header.lattice[2]))
if determinant >= 0:
return destination
for index in range(2, 5):
values = [-float(item) for item in lines[index].split()]
lines[index] = " ".join(f"{value:.16g}" for value in values)
_write_text_atomic(destination, "\n".join(lines) + "\n")
return destination
[docs]
def scale_poscar_lattice(
factor: float,
path: str | os.PathLike[str] = "POSCAR",
) -> Path:
"""Multiply POSCAR's universal linear scale by a positive factor."""
if not math.isfinite(factor) or factor <= 0:
raise ValueError("lattice scale factor must be positive and finite")
destination = Path(path)
lines = destination.read_text(encoding="utf-8").splitlines()
if len(lines) < 2:
raise ValueError("POSCAR has no scale line")
try:
scale = float(lines[1].split()[0])
except (IndexError, ValueError) as exc:
raise ValueError("POSCAR has an invalid scale line") from exc
lines[1] = f"{scale * factor:.16g}"
_write_text_atomic(destination, "\n".join(lines) + "\n")
return destination
[docs]
def derive_seed(entropy: str) -> int:
"""Derive one reproducible 63-bit seed from a caller-supplied string.
The string is the caller's own identity of the perturbation — an attempt
ordinal, a job key, a remedy step — so the same attempt always derives the
same seed and two different attempts derive different ones, without anything
reading a clock or a global random state.
"""
if not isinstance(entropy, str) or not entropy:
raise ValueError("seed entropy must be a nonempty string")
return int.from_bytes(hashlib.sha256(entropy.encode("utf-8")).digest()[:8], "big") >> 1
[docs]
def rattle_poscar(
path: str | os.PathLike[str] = "POSCAR",
*,
amplitude: float = 0.01,
seed: int | None = None,
entropy: str | None = None,
) -> Path:
"""Apply a deterministic bounded perturbation to POSCAR site coordinates.
A perturbation has to be reproducible *and* different between two attempts of
the same calculation, so this function never invents entropy of its own:
either *seed* names the stream explicitly, or *entropy* is a string
:func:`derive_seed` turns into one — typically something attempt-derived, such
as ``f"{job_key}:{attempt_ordinal}"``. Giving neither is refused rather than
silently repeated, because two retries that rattle identically are two
identical calculations.
"""
if not math.isfinite(amplitude) or amplitude < 0:
raise ValueError("rattle amplitude must be finite and nonnegative")
if seed is None and entropy is None:
raise ValueError(
"rattle_poscar needs an explicit seed or an entropy string to derive one from; "
"pass entropy=f'{job_key}:{attempt_ordinal}' to make every retry differ reproducibly"
)
if entropy is not None:
seed = derive_seed(entropy if seed is None else f"{seed}:{entropy}")
destination = Path(path)
lines = destination.read_text(encoding="utf-8").splitlines()
header = read_poscar_header(destination)
coordinate_start = 7
if len(lines) > coordinate_start and lines[coordinate_start].strip().lower().startswith("s"):
coordinate_start += 1
if len(lines) <= coordinate_start:
raise ValueError("POSCAR has no coordinate mode")
coordinate_start += 1
total = sum(header.counts)
if len(lines) < coordinate_start + total:
raise ValueError("POSCAR has fewer coordinate rows than declared atoms")
generator = random.Random(seed)
for index in range(coordinate_start, coordinate_start + total):
fields = lines[index].split()
if len(fields) < 3:
raise ValueError(f"invalid POSCAR coordinate row {index + 1}")
values = [float(fields[column]) + generator.uniform(-amplitude, amplitude) for column in range(3)]
lines[index] = " ".join([*(f"{value:.14f}" for value in values), *fields[3:]])
_write_text_atomic(destination, "\n".join(lines) + "\n")
return destination
def _write_text_atomic(path: Path, value: str) -> None:
descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
temporary = Path(temporary_name)
try:
with os.fdopen(descriptor, "w", encoding="utf-8") as stream:
stream.write(value)
os.replace(temporary, path)
finally:
temporary.unlink(missing_ok=True)
[docs]
def calculate_nbands(
*,
poscar: str | os.PathLike[str] = "POSCAR",
potcar: str | os.PathLike[str] = "POTCAR",
incar: str | os.PathLike[str] = "INCAR",
divisor: int | None = None,
) -> int:
"""Calculate a conservative VASP band count from input metadata.
The count is a heuristic, deliberately generous margin over VASP's own
default of roughly half the valence electrons plus half the atom count: too
few bands is a run that stops with ``TOO FEW BANDS`` or converges to the wrong
state, while a handful of extra empty bands costs only time. The spin-polarized
branch adds ``0.6 * electrons`` plus a magnetization- or atom-count-derived
margin, and the maximum of the candidate formulas wins; the numbers themselves
are empirical defaults carried over from *httk* v1. The result is rounded up to
an even number, and to a multiple of *divisor* when a parallel band divisor
(``NPAR``) is in play, because VASP would otherwise round NBANDS up itself and
report a changed value.
"""
if divisor is not None and (isinstance(divisor, bool) or divisor < 1):
raise ValueError("NBANDS divisor must be a positive integer")
header = read_poscar_header(poscar)
text = Path(potcar).read_text(encoding="utf-8", errors="replace")
zvals = [float(value) for value in re.findall(r"\bZVAL\s*=\s*([-+0-9.Ee]+)", text)]
if len(zvals) < len(header.counts):
raise ValueError("POTCAR contains fewer ZVAL entries than POSCAR species")
electrons = sum(count * zval for count, zval in zip(header.counts, zvals, strict=False))
tags = read_incar(incar)
ispin = int(float(tags.get("ISPIN", "1")))
natoms = max(6, sum(header.counts))
candidates: tuple[int, ...]
if ispin == 2:
magnetic = _expanded_sum(tags.get("MAGMOM", suggested_magnetic_moments(poscar)))
candidates = (
math.floor(0.6 * electrons + 1) + math.ceil(natoms / 2),
math.floor(0.6 * electrons + 1) + math.ceil(abs(magnetic) / 2),
math.floor(0.6 * electrons + 1) + 20,
)
else:
candidates = (
math.floor(electrons / 2 + 2) + math.ceil(natoms / 2),
math.ceil(electrons / 2) + 20,
)
result = max(candidates)
if result % 2:
result += 1
if divisor is not None and result % divisor:
result += divisor - result % divisor
return result
def _expanded_sum(value: str) -> float:
total = 0.0
for item in value.split():
if "*" in item:
count, number = item.split("*", 1)
total += int(count) * float(number)
else:
total += float(item)
return total
[docs]
def potcar_summary(
path: str | os.PathLike[str] = "POTCAR",
output: str | os.PathLike[str] = "POTCAR.summary",
) -> Path:
"""Write a non-potential metadata summary suitable for logs."""
text = Path(path).read_text(encoding="utf-8", errors="replace")
blocks: list[str] = []
for block in re.split(r"(?=^\s*TITEL\s*=)", text, flags=re.MULTILINE):
if "TITEL" not in block:
continue
keep: list[str] = []
for line in block.splitlines():
if re.match(r"^\s*(TITEL|POMASS|ZVAL|ENMAX|ENMIN|LEXCH|EATOM)\b", line):
keep.append(line.rstrip())
if keep:
blocks.append("\n".join(keep))
destination = Path(output)
_write_text_atomic(destination, "\n\n".join(blocks) + ("\n" if blocks else ""))
return destination
[docs]
@dataclass(frozen=True)
class VaspPreparationOptions:
"""Options for dependency-free VASP input preparation.
``accuracy_per_atom`` is the target total-energy accuracy in eV per atom:
``EDIFFG`` is set to that budget for the whole cell and ``EDIFF`` to the same
budget divided by ``ediff_margin``, so the electronic loop converges roughly
one and a half orders of magnitude tighter than the ionic loop it feeds. The
margin of ``33`` is a heuristic, empirical default carried over from *httk*
v1, not a derived quantity.
"""
kpoint_density: float = 20.0
centering: str = DEFAULT_KPOINT_CENTERING
accuracy_per_atom: float | None = 0.001
ediff_margin: float = 33.0
pseudopotential_library: str | os.PathLike[str] | None = None
parallel_tag: str | None = None
parallel_value: int | None = None
normalize_handedness: bool = True
incar_tags: Mapping[str, object] = field(default_factory=dict)