from __future__ import annotations
import numbers
import os
from abc import abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING
if TYPE_CHECKING:
import pathlib
from types import NotImplementedType
from typing import Literal, Self
import ase.io
import numpy as np
from scipy import optimize
from directupsampling.data import GPa_to_eV_per_Angstrom3
[docs]
@dataclass
class RawData:
"""Container for raw data used in EOS fitting."""
volume: np.ndarray
energy: np.ndarray
[docs]
def check_pressure_unit(unit: str) -> float:
"""Return the conversion factor from the given pressure unit to eV/ų.
Parameters
----------
unit : str
Pressure unit string. Accepted values (case-insensitive, spaces and
underscores ignored): ``"eV_per_Angstrom3"`` and ``"GPa"``.
Returns
-------
float
Multiplicative factor to convert a value in ``unit`` to eV/ų.
Raises
------
ValueError
If ``unit`` is not recognised.
"""
unit = unit.lower()
unit = unit.replace("_", "").replace(" ", "")
unit = unit.replace("strom", "")
if unit not in ["evperang3", "gpa"]:
msg = "pressure unit must be 'GPa' or 'eV_per_Angstrom3'"
raise ValueError(msg)
return GPa_to_eV_per_Angstrom3 if unit == "gpa" else 1
[docs]
class EOS:
def __init__(
self,
eq_energy: float,
eq_volume: float,
eq_bulk_modulus: float,
eq_bulk_modulus_derivative: float,
bulk_modulus_unit: Literal["eV_per_Angstrom3", "GPa"] = "eV_per_Angstrom3",
name: str = "EOS",
):
"""Initiate EOS.
Parameters
----------
eq_energy: float
Energy at the equilibrium volume.
eq_volume: float
Equilibrium (0 pressure) volume.
eq_bulk_modulus: float
Bulk modulus at equilibrium volume.
eq_bulk_modulus_derivative: float
Bulk modulus derivative at equilibrium volume.
bulk_modulus_unit: {"GPa", "eV_per_Angstrom3"}
The unit in which the bulk modulus is given. The default is "eV_per_Angstrom3".
name: str
Note
----
The eq_bulk_modulus is always stored as eV/Ang^3 internally.
"""
unit_conversion = check_pressure_unit(bulk_modulus_unit)
self.eq_energy = eq_energy
self.eq_volume = eq_volume
self.eq_bulk_mod = eq_bulk_modulus * unit_conversion
self.eq_bulk_mod_deriv = eq_bulk_modulus_derivative
self.name = name
self.raw_data = None
self.volume_bounds = np.array([0.92, 1.12]) * self.eq_volume
[docs]
@abstractmethod
def get_energy(self, volume: float) -> float: ...
[docs]
@abstractmethod
def get_pressure(self, volume: float, unit: str) -> float: ...
[docs]
@abstractmethod
def get_bulk_modulus(self, volume: float, unit: str) -> float: ...
def __call__(self, volume: float) -> float:
return self.get_energy(volume)
def __add__(self, b: object) -> EOS | NotImplementedType:
if isinstance(b, numbers.Real):
a = self.asarray()
a[0] += b
return Vinet(*a, bulk_modulus_unit="evperang3")
return NotImplemented
def __radd__(self, a: object) -> EOS | NotImplementedType:
return self.__add__(a)
[docs]
def get_volume(self, pressure: float, unit: str = "eV_per_Angstrom3") -> float:
"""Find the volume at the given pressure by scalar optimisation.
Parameters
----------
pressure : float
Target pressure.
unit : {"eV_per_Angstrom3", "GPa"}
Unit of ``pressure``. Default is eV/ų.
Returns
-------
float
Volume in ų at which the EOS pressure equals ``pressure``.
"""
unit_conversion = check_pressure_unit(unit)
pressure *= unit_conversion
def delta_p2(volume):
return (self.get_pressure(volume) - pressure) ** 2
# bracket: just some v>0 and very large one to avoid div by 0 and inf
bbracket = [1e-15, self.eq_volume]
result = optimize.minimize_scalar(delta_p2, bracket=bbracket, tol=1e-10)
opt_volume = result.x
# Assert that the optimized volume gives the right pressure (+- 100 Pa)
np.testing.assert_allclose(
self.get_pressure(opt_volume), pressure, atol=1e-7, rtol=0
)
return opt_volume
[docs]
def todict(self) -> dict:
"""Return a JSON-serialisable dict of the EOS parameters.
Returns
-------
dict
Dictionary of EOS parameters suitable for ``write_json``.
"""
return {
"eq_energy": self.eq_energy,
"eq_volume": self.eq_volume,
"eq_bulk_mod": self.eq_bulk_mod,
"eq_bulk_mod_deriv": self.eq_bulk_mod_deriv,
"name": self.name,
"volume_bounds": self.volume_bounds,
}
[docs]
@classmethod
def fromdict(cls, dct: dict) -> Self:
"""Reconstruct an EOS instance from a dict produced by ``todict``.
Parameters
----------
dct : dict
Dictionary as returned by :meth:`todict`.
Returns
-------
Self
New instance with parameters restored from ``dct``.
"""
new = cls(
eq_energy=dct["eq_energy"],
eq_volume=dct["eq_volume"],
eq_bulk_modulus=dct["eq_bulk_mod"],
eq_bulk_modulus_derivative=dct["eq_bulk_mod_deriv"],
bulk_modulus_unit="eV_per_Angstrom3",
name=dct.get("name", "EOS"),
)
new.volume_bounds = dct["volume_bounds"]
return new
[docs]
def asarray(self) -> np.ndarray:
"""Return the four EOS parameters as a 1-D array in internal units.
Returns
-------
np.ndarray
``[eq_energy, eq_volume, eq_bulk_mod, eq_bulk_mod_deriv]`` in
eV and ų.
"""
return np.array(
[self.eq_energy, self.eq_volume, self.eq_bulk_mod, self.eq_bulk_mod_deriv]
)
[docs]
def write(self, filename: str | pathlib.Path) -> None:
"""Write the EOS parameters to a text file in the legacy format.
Parameters
----------
filename : str or pathlib.Path
Output file path.
"""
coef = self.asarray()
# Write according to old format in meV and GPa
coef[0] *= 1e3
coef[2] /= GPa_to_eV_per_Angstrom3
np.savetxt(filename, coef[np.newaxis, :], fmt="%.14f")
[docs]
@classmethod
def read(cls, filename: str | pathlib.Path) -> Self:
"""Read EOS parameters from a legacy text file written by :meth:`write`.
Parameters
----------
filename : str or pathlib.Path
Path to the file to read.
Returns
-------
Self
New instance constructed from the file contents.
"""
filename = os.path.expanduser(filename)
coef = np.loadtxt(filename)
coef[0] = coef[0] / 1e3 # Read to eV
coef[2] *= GPa_to_eV_per_Angstrom3
eos = cls(*coef)
return eos
[docs]
@classmethod
def fit(cls, volumes: np.ndarray, energies: np.ndarray) -> Self:
"""Fit an EOS to energy-volume data and return the optimised instance.
Parameters
----------
volumes : np.ndarray
Volumes in ų/atom.
energies : np.ndarray
DFT total energies in eV/atom, same length as ``volumes``.
Returns
-------
Self
Fitted EOS instance with ``raw_data`` and ``volume_bounds`` set.
"""
v = np.array(volumes, dtype=float).ravel()
e = np.array(energies, dtype=float).ravel()
tmp = cls(*(1,) * 4)
def fit_func(vol, e0, v0, b0, bp):
tmp.eq_energy = e0
tmp.eq_volume = v0
tmp.eq_bulk_mod = b0
tmp.eq_bulk_mod_deriv = bp
return tmp.get_energy(vol)
iemin = np.argmin(e)
emin, v_emin = e[iemin], v[iemin]
p0 = [emin, v_emin, 1e-3, 4]
bounds = ([-np.inf, 0, 0, 0], [np.inf, np.inf, np.inf, np.inf])
result = optimize.curve_fit(fit_func, v, e, p0=p0, bounds=bounds)
new = cls(*result[0])
new.raw_data = RawData(v, e)
new.set_bounds_from_raw_data()
return new
[docs]
def set_bounds_from_raw_data(self) -> None:
"""Set ``volume_bounds`` from the volume range in ``raw_data``."""
vmin = self.raw_data.volume.min()
vmax = self.raw_data.volume.max()
if vmin - 0.2 * (vmax - vmin) > 0:
self.volume_bounds[0] = vmin - 0.2 * (vmax - vmin)
else:
self.volume_bounds[0] = 0
self.volume_bounds[1] = vmax + 0.2 * (vmax - vmin)
[docs]
class Vinet(EOS):
directupsampling_objtype = "vinet"
[docs]
def get_energy(self, volume: float) -> float:
"""Return the Vinet energy at the given volume.
Parameters
----------
volume : float
Volume in ų.
Returns
-------
float
Energy in eV.
"""
e0 = self.eq_energy
v0 = self.eq_volume
bm0 = self.eq_bulk_mod
bmd0 = self.eq_bulk_mod_deriv
eta = (volume / v0) ** (1 / 3)
factor = 2 * bm0 * v0 * (bmd0 - 1) ** -2
exponent = 3 / 2 * (bmd0 - 1) * (1 - eta)
energy = e0 + factor * (
2 - (5 + 3 * bmd0 * (eta - 1) - 3 * eta) * np.exp(exponent)
)
return energy
[docs]
def get_pressure(
self,
volume: float,
unit: Literal["eV_per_Angstrom3", "GPa"] = "eV_per_Angstrom3",
) -> float:
"""Get the analytical pressure at the specified volume.
Parameters
----------
volume : float
The volume at which to calculate the pressure.
unit : {"eV_per_Angstrom3", "GPa"}
Specifies the unit of the returned pressure. The default is eV/ų.
Returns
-------
float
Pressure in the requested unit.
"""
unit_conversion = 1 / check_pressure_unit(unit)
v0 = self.eq_volume
bm0 = self.eq_bulk_mod
bmd0 = self.eq_bulk_mod_deriv
eta = (volume / v0) ** (1 / 3)
factor = 3 * bm0 * ((1 - eta) / eta**2)
exponent = -3 / 2 * (bmd0 - 1) * (eta - 1)
pressure = factor * np.exp(exponent)
return pressure * unit_conversion
[docs]
def get_bulk_modulus(
self,
volume: float,
unit: Literal["eV_per_Angstrom3", "GPa"] = "eV_per_Angstrom3",
) -> float:
"""Get the analytical bulk modulus at the specified volume.
Parameters
----------
volume : float
The volume at which to calculate the bulk modulus.
unit : {"eV_per_Angstrom3", "GPa"}
Specifies the unit of the returned bulk modulus. The default is eV/ų.
Returns
-------
float
Bulk modulus in the requested unit.
"""
unit_conversion = 1 / check_pressure_unit(unit)
v0 = self.eq_volume
bm0 = self.eq_bulk_mod
bmd0 = self.eq_bulk_mod_deriv
eta = (volume / v0) ** (1 / 3)
factor = 3 * bm0 * ((1 - eta) / eta**2)
exponent = -3 / 2 * (bmd0 - 1) * (eta - 1)
deta = 1 / 3 * (volume / v0) ** (-2 / 3) / v0
dexp = -3 / 2 * (bmd0 - 1) * deta
dfac = -deta * eta**-2 - 2 * (1 - eta) * eta**-3 * deta
dpdv = 3 * bm0 * np.exp(exponent) * (dfac + factor * dexp)
bm = -volume * dpdv
return bm * unit_conversion
[docs]
def read_eos(file: pathlib.Path) -> EOS:
"""Read an EOS from a file, dispatching on the filename.
Parameters
----------
file : pathlib.Path
Path to the EOS file. The filename must contain ``"vinet"``
(case-insensitive) to be recognised as a Vinet EOS.
Returns
-------
EOS
Parsed EOS instance.
Raises
------
RuntimeError
If the filename does not match any known EOS type.
"""
name = file.name
if "vinet" in name.lower():
return Vinet.read(file)
msg = f"Don't know how to read {name} ({file})"
raise RuntimeError(msg)
[docs]
def get_eos_from_folders(paths: list[pathlib.Path]) -> EOS:
"""Fit a Vinet EOS from a list of completed VASP calculation directories.
Parameters
----------
paths : list of pathlib.Path
Directories each containing a ``vasprun.xml`` file. Each directory
represents one volume point.
Returns
-------
EOS
Fitted Vinet EOS instance.
"""
volumes = []
energies = []
for calc_dir in paths:
atoms = ase.io.read(calc_dir / "vasprun.xml")
natoms = len(atoms)
volumes.append(atoms.get_volume() / natoms)
energies.append(atoms.get_potential_energy() / natoms)
return Vinet.fit(volumes, energies)