from __future__ import annotations
import logging
import os
import pathlib
import re
import warnings
from typing import TYPE_CHECKING
from xml.etree import ElementTree
import ase.io
import numpy as np
import scipy.constants
import scipy.special
from ase.atoms import Atoms
from ase.calculators.calculator import ReadError
from ase.calculators.singlepoint import SinglePointDFTCalculator
from directupsampling.parallel import paropen
from directupsampling.utils import get_calcs_names
if TYPE_CHECKING:
from directupsampling.sampling import GridSampler
logger = logging.getLogger(__name__)
KB = scipy.constants.physical_constants["Boltzmann constant in eV/K"][0]
# Don't support OUTCAR anymore, since no parameters are read
DEFAULT_RESULTS_FILES = ["vasprun.xml"] # , "OUTCAR"]
[docs]
def read_grid_from_folders(path: os.PathLike) -> list[tuple[float, float]]:
path = pathlib.Path(path)
if not path.is_dir():
raise FileNotFoundError(f"Directory {path} does not exist.")
dirs = sorted(path.glob("*Ang3_*K*/"))
if len(dirs) == 0:
grid = _read_grid_from_latpar_folders(path)
else:
grid = []
for d in dirs:
volume = re.findall(r"([0-9.]+)Ang3", str(d))[0]
temper = re.findall(r"([0-9.]+)K", str(d))[0]
gp = (volume, temper)
grid.append(gp)
return grid
def _read_grid_from_latpar_folders(path: pathlib.Path) -> list[tuple[float, float]]:
dirs = sorted(path.glob("*Ang_*K*/"))
grid = []
for d in dirs:
volume = _get_volume_from_calc(d)
temper = re.findall(r"([0-9.]+)K", str(d))[0]
gp = (volume, temper)
grid.append(gp)
return grid
def _get_volume_from_calc(path: pathlib.Path) -> float:
dirs = sorted(path.glob("*step*/"))
if len(dirs) == 0:
dirs = sorted(path.glob("*/*step*/"))
for d in dirs:
try:
atoms: Atoms = _read_calc_dir(d)["atoms"]
return atoms.get_volume() / len(atoms)
except (FileNotFoundError, ElementTree.ParseError, ReadError):
continue
else:
raise RuntimeError(f"No complete calculations found in {path}.")
[docs]
def read_folders_to_sampler(sampler: GridSampler, path: os.PathLike, **kwargs):
root_path = pathlib.Path(path)
for gp in sampler.grid:
grid_dirname = f"{gp.volume}Ang3_{gp.temperature}K"
path_up_to_grid = root_path / grid_dirname
if not path_up_to_grid.is_dir():
path_up_to_grid = _find_dirname_from_calcs(gp, root_path)
_iread(sampler, path_up_to_grid, **kwargs)
def _find_dirname_from_calcs(gp, root_path: pathlib.Path):
dirs = sorted(root_path.glob("*_*K*/"))
for d in dirs:
volume = _get_volume_from_calc(d)
temper = float(re.findall(r"([0-9.]+)K", str(d))[0])
if (
abs(volume - float(gp.volume)) < 1e-7
and abs(temper - float(gp.temperature)) < 1e-4
):
path_up_to_grid = d
break
else:
raise RuntimeError(f"No results folders found in {root_path}.")
return path_up_to_grid
def _iread(sampler: GridSampler, path_up_to_grid: pathlib.Path, **kwargs: dict) -> None:
dirs = sorted(path_up_to_grid.glob("*step*/"))
if len(dirs) == 0:
dirs = sorted(path_up_to_grid.glob("*/*step*/"))
for directory in dirs:
seed_matches = re.findall(r"seed_*([0-9]+)", str(directory))
seed = float(seed_matches[0]) if len(seed_matches) > 0 else 0
step = float(re.findall(r"step_*([0-9]+)", str(directory))[0])
gp = sampler.grid.current
index = tuple(getattr(gp, ax) for ax in sampler.grid.axes) + (seed, step)
try:
dct = _read_calc_dir(directory, **kwargs)
except (FileNotFoundError, ElementTree.ParseError, ReadError):
continue
_check_atoms_and_calc(dct["atoms"], float(index[0]), float(index[1]))
if isinstance(dct["atoms"].calc, SinglePointDFTCalculator):
dct["electronic_free_energy"] = get_fixdos_fel(
dct["atoms"].calc, float(index[1])
)
sampler.snapshot_container.insert(index, **dct)
def _read_calc_dir(directory: pathlib.Path, old_fermi_folders=False) -> dict:
try:
atoms = _read_calc_output(directory)
except FileNotFoundError as e:
# This is only to be able to read old folders
try:
subdir = "calc_kBT" if old_fermi_folders else "calc_0K"
atoms = _read_calc_output(directory / subdir)
except FileNotFoundError:
raise e # raise the previous "up-to-date" error
ref_names, ref_energies = _read_ref_energy_files(directory)
names = get_calcs_names(*ref_names, atoms.calc.name)
dct = {}
for name, energy in zip(names, ref_energies):
dct[f"potential_energy_{name}"] = energy
dct["atoms"] = atoms
return dct
def _read_calc_output(directory: pathlib.Path):
results_files = DEFAULT_RESULTS_FILES
for results_file in results_files:
for file in directory.iterdir():
if file.stem == pathlib.Path(results_file).stem:
calc_file = file
break
else:
continue
break
else:
raise FileNotFoundError(str(directory) + str(results_files) + "[.gz]")
msg = f"The file {calc_file} is incomplete and could not be read properly."
try:
# parallel=False to circumvent a bug (see issue #1431) in ASE that results
# in a parallel deadlock when reading in parallel when a POSCAR is present
# in the same directory as the OUTCAR. Update if ASE MR !3253 goes into
# latest release.
atoms = Atoms(ase.io.read(calc_file, parallel=False))
except ElementTree.ParseError as e:
raise ReadError(msg) from e
if atoms.calc is None:
raise ReadError(msg)
return atoms
def _read_ref_energy_files(directory: pathlib.Path):
names = []
energies = []
energy_files = sorted(directory.glob("energy_*"))
for energy_file in energy_files:
names.append(str(energy_file.name).replace("energy_", "", 1))
with paropen(energy_file, "r") as f:
energies.append(float(f.readlines()[0]))
return names, energies
def _check_atoms_and_calc(atoms: Atoms, volume: float, temperature: float):
volume_per_atom = atoms.get_volume() / len(atoms)
volume_from_grid = float(volume)
np.testing.assert_allclose(
volume_per_atom,
volume_from_grid,
err_msg=f"The volume read from calculation output ({volume_per_atom}) "
f"does not equal the intended volume from the grid ({volume_from_grid}).",
)
if atoms.calc.parameters.get("ismear") == -1:
sigma = atoms.calc.parameters["sigma"]
np.testing.assert_allclose(
sigma,
temperature * KB,
atol=1e-6,
err_msg=f"The sigma read from calculation output ({sigma}) does not "
f"correspond to the temperature * 'Boltzmann constant' "
f"({temperature * KB}).",
)
[docs]
def get_fixdos_fel(calc: SinglePointDFTCalculator, temperature: float) -> float:
"""Calculate the electronic free energy within the fixed DOS approximation.
Parameters
----------
calc: ASE SinglePointDFTCalculator
Calculator instance containing eigenvalues.
temperature: float
The 'electronic' temperature going into the Fermi Dicac distribution.
Returns
-------
free_energy
The free energy at `temperature`.
"""
import scipy.optimize
beta = 1.0 / (temperature * KB)
w_k = calc.get_k_point_weights()
e_kn = np.array([calc.get_eigenvalues(kpt=k) for k in range(len(w_k))])
w_kn = np.repeat(w_k[:, np.newaxis], e_kn.shape[1], axis=1)
fermi_energy = calc.get_fermi_level()
where_occupied = e_kn <= fermi_energy
lumo_zero_kelvin = np.min(e_kn[~where_occupied])
integrated_nel_zero_kelvin = 2 * np.sum(w_kn[where_occupied])
integrated_ene_zero_kevlin = 2 * np.sum(e_kn[where_occupied] * w_kn[where_occupied])
def number_of_electrons(fermi_level):
return 2 * np.sum(fermi_dirac_distribution(e_kn, fermi_level, beta) * w_kn)
fermi_level = scipy.optimize.minimize_scalar(
lambda _: (number_of_electrons(_) - integrated_nel_zero_kelvin) ** 2,
bracket=[np.min(e_kn), np.max(e_kn)],
)["x"]
internal_energy = (
2 * np.sum(fermi_dirac_distribution(e_kn, fermi_level, beta) * w_kn * e_kn)
- integrated_ene_zero_kevlin
)
f = fermi_dirac_distribution(e_kn, fermi_level, beta)
where_partial = (0 < f) & (f < 1)
f = f[where_partial]
entropy = -2 * np.sum(
w_kn[where_partial] * (f * np.log(f) + (1 - f) * np.log(1 - f))
)
free_energy = internal_energy - entropy / beta
results = {
"number of electrons": integrated_nel_zero_kelvin,
"fermi shift": fermi_level - lumo_zero_kelvin, # eV
"occupancy at highest level": fermi_dirac_distribution(
np.max(e_kn), fermi_level, beta
),
# or should be - fermi_energy? How does VASP calculate fermi_energy?
"internal energy": internal_energy, # eV
"entropy": entropy, # kB
"free energy": free_energy, # eV
}
# 1e-4 is the precision written in vasprun.xml
if results["occupancy at highest level"] > 1e-4:
warnings.warn(
f"The highest band level is occupied: "
f"{results['occupancy at highest level']}"
)
return results["free energy"]
[docs]
def fermi_dirac_distribution(
energy: float | np.ndarray,
fermi_level: float,
beta: float,
) -> np.ndarray:
"""Compute the Fermi–Dirac occupation probability for each energy level.
Parameters
----------
energy : ndarray of shape (...,), dtype=float64
Energy values.
fermi_level : float
The Fermi level.
beta : float
Inverse temperature 1 / (k_B * T).
Returns
-------
ndarray of shape (...,), dtype=float64
The occupation probability for each energy level.
"""
# expit replaces 1 / (exp(-x) + 1), but is numerically safer
return scipy.special.expit(-(energy - fermi_level) * beta)