from __future__ import annotations
import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING, Protocol
import ase.io
import numpy as np
import scipy.constants
try:
from ase.calculators.vasp import Vasp as _VaspCalc
from ase.calculators.vasp.create_input import open_potcar as _open_vasp_potcar
except ImportError:
_VaspCalc = None
_open_vasp_potcar = None
from directupsampling.parallel import parcall, paropen
from directupsampling.parallel import world as comm
from .read import _read_calc_dir
if TYPE_CHECKING:
from ase.atoms import Atoms
from ase.calculators.calculator import FileIOCalculator
from directupsampling.sampling.sampling import GridSampler
from directupsampling.snapshots import SnapshotContainer
logger = logging.getLogger(__name__)
KB = scipy.constants.physical_constants["Boltzmann constant in eV/K"][0]
NBANDS_FACTOR_MIN = 0.95
NBANDS_FACTOR_MAX = 1.1
NBANDS_MAXTEMP = 3950
[docs]
def write_folders_from_sampler(
sampler: GridSampler,
calc: FileIOCalculator,
root_path: str | os.PathLike | None = None,
**kwargs: dict,
) -> tuple[list, list]:
"""Write folder structure from a SnapshotSampler instance.
Returns
-------
joblist : list
joblist_angk : list
"""
root_path = Path(root_path) if root_path is not None else Path.cwd()
joblist = []
joblist_angk = []
grid = sampler.grid
parcall(logger.info, "Writing folders for %s with %s.", grid, calc, comm=comm)
fmt = kwargs.pop("fmt", None)
for gp in grid:
grid_dirname = Path(f"{gp.volume}Ang3_{gp.temperature}K")
path = root_path / grid_dirname
if calc:
ref_atoms = sampler.atoms or sampler.snapshot_container.get("atoms")[0]
_check_calc(calc, ref_atoms, gp)
snapshots = sampler.snapshot_container.filter(
f"volume == {gp.volume} & temperature == {gp.temperature}",
)
subjoblist = write_snapshots(path, snapshots, calc, sampler, fmt=fmt)
_write_ref_file(
path / "natoms", len(sampler.snapshot_container.get("atoms")[0])
)
_write_ref_file(path / "temp", gp.temperature)
joblist.extend(subjoblist)
joblist_angk.append(path)
n = len(joblist)
parcall(logger.info, "Wrote %d director%s.", n, "y" if n == 1 else "ies", comm=comm)
return joblist, joblist_angk
[docs]
def write_snapshots(
path: Path,
snapshots: SnapshotContainer,
calc: FileIOCalculator = None,
sampler: GridSampler | None = None,
fmt: str | None = None,
) -> list[Path]:
"""Write snapshots to folders.
Parameters
----------
path : Path
Root directory for this set of snapshots.
snapshots : SnapshotContainer
Snapshots to write.
calc : FileIOCalculator | None
ASE FileIOCalculator to write calculator-specific input files.
If None, a structure file is written using *fmt*.
sampler : GridSampler | None
Used to look up reference energies when atoms carry no calculator.
fmt : str | None
ASE format string for the structure file when *calc* is None.
Defaults to ``"vasp"`` (writes a ``POSCAR`` file).
Returns
-------
joblist : list[Path]
"""
if fmt is None:
fmt = "vasp"
outname = "POSCAR" if fmt == "vasp" else f"structure.{fmt}"
joblist = []
for index, atoms in zip(snapshots.index, snapshots.get("atoms"), strict=True):
atoms: Atoms
remaining_index_dirname = Path(f"seed{int(index.seed)}_step{int(index.step)}")
dirname = path / remaining_index_dirname
if atoms.calc is not None:
ref_energy = atoms.get_potential_energy()
ref_name = atoms.calc.name
else:
ref_energy = sampler.snapshot_container.get(index, "potential_energy")
ref_name = "ref"
if ref_energy is not None:
_write_ref_file(dirname / Path("energy_" + ref_name), ref_energy)
if _is_calc_done(dirname):
parcall(
logger.warning,
"%s already contains calculation. "
"Only reference energy was added/updated.",
dirname,
comm=comm,
)
continue
if calc:
calc.directory = dirname
if comm.rank == 0:
calc.write_input(atoms)
else:
if comm.rank == 0:
dirname.mkdir(parents=True, exist_ok=True)
ase.io.write(dirname / outname, atoms, format=fmt)
joblist.append(dirname)
return joblist
class _GridPoint(Protocol):
temperature: float
volume: float
def _param_get(params: dict, key: str) -> object | None:
"""Case-insensitive lookup in a VASP parameter dict.
Returns
-------
object | None
The value for the matching key, or None if not found.
"""
key_lower = key.lower()
for k, v in params.items():
if k.lower() == key_lower:
return v
return None
def _check_calc(calc: FileIOCalculator, atoms: Atoms, gp: _GridPoint) -> None:
_try_vasp_initialize(calc, atoms)
if not _param_get(calc.parameters, "encut"):
enmax = _read_enmax(calc)
if enmax:
calc.set(encut=enmax)
if _param_get(calc.parameters, "ismear") == -1:
calc.set(sigma=gp.temperature * KB)
if not _param_get(calc.parameters, "nbands"):
nelect = _read_nelect(calc)
if nelect is not None:
natoms = len(atoms)
nbands_default = int(
np.max(
[
np.round(natoms / 2 + nelect / 2),
np.round(nelect * 0.6),
]
)
)
nbands = int(
np.round(
nbands_default
* (
NBANDS_FACTOR_MIN
+ (NBANDS_FACTOR_MAX - NBANDS_FACTOR_MIN)
* gp.temperature
/ NBANDS_MAXTEMP
)
)
)
calc.set(nbands=nbands)
def _try_vasp_initialize(calc: FileIOCalculator, atoms: Atoms) -> None:
"""Call Vasp.initialize to populate ppp_list from VASP_PP_PATH.
write_input calls initialize anyway, so calling it here is idempotent.
Skips silently if calc is not a Vasp instance or VASP_PP_PATH is unset.
"""
if _VaspCalc is None or not isinstance(calc, _VaspCalc):
return
try:
calc.initialize(atoms)
except (RuntimeError, FileNotFoundError, OSError):
logger.debug("Vasp.initialize failed; POTCAR-based defaults unavailable.")
def _read_enmax(calc: FileIOCalculator) -> float | None:
ppp_list = getattr(calc, "ppp_list", None)
if not ppp_list or _open_vasp_potcar is None:
return None
enmax = 0.0
try:
for potcar_path in ppp_list:
with _open_vasp_potcar(potcar_path) as f:
for line in f:
if "ENMAX" in line:
enmax = max(enmax, float(line.split()[2].replace(";", "")))
except OSError:
logger.debug("Failed to read ENMAX from source POTCARs.")
return None
return enmax or None
def _read_nelect(calc: FileIOCalculator) -> int | None:
if (nelect := _param_get(calc.parameters, "nelect")) is not None:
return nelect
has_ppp = hasattr(calc, "default_nelect_from_ppp") and getattr(
calc, "ppp_list", None
)
if not has_ppp:
return None
try:
return int(calc.default_nelect_from_ppp())
except (AssertionError, RuntimeError, OSError):
logger.debug("default_nelect_from_ppp failed.")
return None
def _is_calc_done(directory: Path) -> bool:
try:
_read_calc_dir(directory)
except (FileNotFoundError, ase.io.ParseError):
is_done = False
else:
is_done = True
return is_done
def _write_ref_file(file: Path, value) -> None:
directory = file.parent
if comm.rank == 0:
directory.mkdir(parents=True, exist_ok=True)
with paropen(file, "w") as f:
f.write(f"{value}\n")
[docs]
def write_joblist(joblist: list[str], filename: str | Path = "jobList") -> None:
file = Path(filename)
if comm.rank == 0:
if file.is_file():
file.rename(file.with_suffix(".old"))
with paropen(file, "a") as f:
for path in joblist:
f.write(f"{Path(path).resolve()}\n")
parcall(logger.info, "Joblist written to %s.", file, comm=comm)