import pathlib
import re
from collections.abc import Callable
from dataclasses import dataclass
import numpy as np
import scipy.optimize
from directupsampling.data import kB
fit_types = ["poly", "fah", "fvib", "fel"]
[docs]
@dataclass
class RawData:
"""Container for raw data used in free energy fitting."""
volume: np.ndarray
temperature: np.ndarray
free_energy: np.ndarray
error: np.ndarray
[docs]
class FreeEnergyContrib:
directupsampling_objtype = "freeenergycontrib"
def __init__(
self,
name: str = "",
bounds: list | None = None,
units: dict = {"volume": "Å$^3$", "temperature": "K", "energy": "eV/atom"},
):
"""Class to hold the free energy data and fit.
Parameters
----------
Notes
-----
The attribute vt_basis is initialized to DEFAULT_VT_BASIS,
but can be set before fitting.
"""
self.name = name
self._bounds = bounds
self.units = units
self.parametrization = None
[docs]
def fit_data(self, volumes, temperatures, free_energies, errors=[], **kwargs):
"""
Parameters
----------
kwargs:
Keyword arguments sent to FreeEnergyFit:
kind: {"poly", "fah", "fvib", "fel"},
vt_basis: str,
mean_freqs: tuple[list[float], list[float]]
"""
v, t, f, e = volumes, temperatures, free_energies, errors
self.parametrization = FreeEnergyFit(v, t, f, e, **kwargs)
self.parametrization.fit()
if self._bounds is None:
self.bounds = _get_bounds_from_data(self.parametrization.raw_data)
@property
def bounds(self):
"""Defaults to [[5, 50], [1, 1500]] if not set manually or by `fit_data()`."""
if self._bounds is None:
return [[5, 50], [1, 1500]]
else:
return self._bounds
@bounds.setter
def bounds(self, bounds):
if np.array(bounds).shape != (2, 2):
raise ValueError("bounds must be 2 by 2")
self._bounds = [list(_) for _ in bounds]
def __repr__(self):
return f"FreeEnergyContrib({self.__dict__!r})"
def __call__(self, volume, temperature):
return self.get_free_energy(volume, temperature)
[docs]
def get_free_energy(self, volume, temperature):
for var, _ in [("volume", volume), ("temperature", temperature)]:
if _ is None:
raise ValueError(f"Argument {var} cannot be None")
return self.parametrization.get_free_energy(volume, temperature)
[docs]
def todict(self):
dct = self.__dict__.copy()
dct["bounds"] = dct.pop("_bounds")
return dct
[docs]
@classmethod
def fromdict(cls, dct):
paramz = dct.pop("parametrization")
new = cls(**dct)
new.parametrization = paramz
return new
[docs]
class FreeEnergyFit:
directupsampling_objtype = "freeenergyfit"
def __init__(
self,
volumes=[],
temperatures=[],
free_energies=[],
errors=[],
kind: str = "poly",
vt_basis: str = (
"1 + v + t + v**2 + vt + t**2"
" + v**3 + v**2*t + v*t**2"
" + v**3*t + v**2*t**2"
),
mean_freqs: tuple[list[float], list[float]] | None = None,
):
self.raw_data = RawData(
np.array(volumes, dtype=float).ravel(),
np.array(temperatures, dtype=float).ravel(),
np.array(free_energies, dtype=float).ravel(),
np.array(errors, dtype=float).ravel(),
)
self.kind = kind
self.vt_basis = VTBasis(vt_basis)
self.mean_freqs = [np.array(_) for _ in mean_freqs] if mean_freqs else None
self.fit_function = None
self.coefficients = None
self.results = {}
[docs]
def get_free_energy(self, volume, temperature):
return self.fit_function((volume, temperature), *self.coefficients)
[docs]
def fit(self):
"""Fit to self.raw_data."""
kwargs = {}
if self.mean_freqs:
kwargs.update({"mean_freqs": self.mean_freqs})
v = self.raw_data.volume
t = self.raw_data.temperature
f = self.raw_data.free_energy
# errors = self.raw_data.errors
v_range = [np.min(v), np.max(v)]
t_range = [np.min(t), np.max(t)]
self.vt_basis.rescale(v_range, t_range)
self.fit_function = get_fit_func(self.vt_basis, self.kind, **kwargs)
guess = 1e-4 * np.ones(len(self.vt_basis(1, 1)))
coeffs, covar = scipy.optimize.curve_fit(
self.fit_function,
(v, t),
f,
p0=guess,
method="trf",
)
self.coefficients = coeffs
residuals = self.get_free_energy(v, t) - f
self.results = {
"covariance": covar,
"residuals": residuals,
"max_residual": np.max(np.absolute(residuals)),
"rmse": np.sqrt(np.mean(residuals**2)),
}
[docs]
def todict(self):
dct = self.__dict__.copy()
del dct["fit_function"]
return dct
[docs]
@classmethod
def fromdict(cls, dct):
raw = dct.pop("raw_data", None)
if isinstance(raw, dict):
new = cls(raw["volume"], raw["temperature"], raw["free_energy"], raw.get("error", []))
elif raw is not None:
new = cls(*raw)
else:
new = cls()
new.__dict__.update(dct)
new.fit_function = get_fit_func(new.vt_basis, new.kind)
return new
def _get_tdict_free_energy(volume: np.ndarray, coeffs: np.ndarray) -> float:
e = 0
for i, coef in enumerate(coeffs):
e += coef * volume**i
return e
def _read_old_fes_format(file: pathlib.Path) -> dict:
tdict = {}
with file.open("r", encoding="utf-8") as fin:
for line in fin:
temperature = int(float(line.split()[0]))
# " * 1e-3": Convert from meV to eV
coef = [float(_) * 1e-3 for _ in line.split()[1:]]
tdict[temperature] = coef
return tdict
def _get_bounds_from_data(data: RawData) -> list[list[float]]:
bounds = []
for a in [data.volume, data.temperature]:
min_a = np.min(a)
max_a = np.max(a)
low_bound = np.max([min_a - 0.2 * (max_a - min_a), 0])
high_bound = max_a + 0.2 * (max_a - min_a)
bounds.append([low_bound, high_bound])
return bounds
[docs]
def get_fit_func(basis: Callable, fit_type: str, **kwargs: dict) -> Callable:
"""Return function for fitting.
Parameters
----------
basis : callable
Callable that takes arguments (volume, temperature) and return a basis
with the length of ``coefficients``.
fit_type : str
One of 'poly', 'fvib', 'fah' and 'fel'. Defaults to 'poly'
kwargs : dict
keyword arguments sent to the respective 'get_fit_function' function:
'fah'; mean_freqs: tuple[np.ndarray, np.ndarray] the arrays contain
volume and frequencies respectively. If mean_freqs is None, 10 meV is
used for all volumes.
Returns
-------
Callable
Raises
------
ValueError
"""
if fit_type == "fvib":
fit_func = get_fvib_fit_func(basis, **kwargs)
elif fit_type == "fah":
fit_func = get_fah_fit_func(basis, **kwargs)
elif fit_type == "poly":
fit_func = get_poly_fit_func(basis, **kwargs)
elif fit_type == "fel":
fit_func = get_fel_fit_func(basis, **kwargs)
else:
raise ValueError(f'fit_type "{fit_type}" not recognized')
return fit_func
[docs]
class VTBasis:
directupsampling_objtype = "vtbasis"
def __init__(self, expression: str, domain=None):
self.expression = expression
self.exponents = self.parse_expression()
self.function = self.build_function()
self.domain = None
if domain is not None:
self.rescale(*domain)
def __call__(self, volume, temperature):
return self.function(volume, temperature)
[docs]
def parse_expression(self):
exponents = []
terms = self.expression.split("+")
for term in terms:
matches = re.findall(r"([a-z])(?:\*\*|\^)?([-.,0-9]*)", term.lower())
pows = {"v": 0, "t": 0} # Only v and t allowed
for match in matches:
pows[match[0]] += float(match[1]) if match[1] else 1
exponents.append(pows)
self.exponents = exponents
return exponents
[docs]
def build_function(self):
def fun(volume, temperature, exponents=self.exponents):
v = np.array(volume)
t = np.array(temperature)
# bshape = 0 * v + 0 * t
basis = []
for pows in exponents:
basis.append(v ** pows["v"] * t ** pows["t"])
return basis
self.function = fun
return self.function
[docs]
def rescale(self, v_range, t_range):
vmin, vmax = v_range
tmin, tmax = t_range
def new_fun(v, t, fun=self.function):
volume_rescaled = -1 + 2 * (v - vmin) / (vmax - vmin) # from -1 to 1
# temp_rescaled = (t - tmin) / (tmax - tmin) # from 0 to 1
temp_rescaled = t / tmax # from 0 to 1
return fun(volume_rescaled, temp_rescaled)
self.function = new_fun
self.domain = [[vmin, vmax], [tmin, tmax]]
[docs]
def todict(self):
dct = {"expression": self.expression, "domain": self.domain}
return dct
def _fvib(
freq: float,
temperature: np.ndarray,
) -> np.ndarray:
"""Give (modified) QM Fvib without ZPE.
Parameters
----------
freq : float
Frequency in eV. Can also be ndarray, broadcastable with `temperature`.
temperature : ndarray of shape (...,), dtype=float64
Temperatures in K.
Returns
-------
ndarray of shape (...,), dtype=float64
Vibrational free energy in eV.
"""
f, t = np.broadcast_arrays(freq, temperature)
valid = (t > 0) & (f > 0)
safe_f = np.where(valid, f, 1.0)
safe_t = np.where(valid, t, 1.0)
result = kB * safe_t * np.log(-np.expm1(-safe_f / (kB * safe_t)))
return np.where(valid, result, 0.0)
[docs]
def get_fah_fit_func(
basis: np.ndarray,
mean_freqs: tuple[list | np.ndarray] | None = None,
mean_freqs_fit_order: int = 2,
):
if mean_freqs is None:
# default to 10 meV
mean_freqs_fit = lambda _: 0.01
else:
mean_freqs_fit = np.polynomial.Polynomial.fit(
np.array(mean_freqs[0]).flat,
np.array(mean_freqs[1]).flat,
mean_freqs_fit_order,
)
def fit_func(vt_tuple, *coef):
v, t = vt_tuple
mean_freq = mean_freqs_fit(v)
return _fvib(
mean_freq + np.dot(np.stack(basis(v, t), axis=-1), coef),
t,
) - _fvib(mean_freq, t)
return fit_func
[docs]
def get_fvib_fit_func(basis):
def fit_func(vt_tuple, *coef, basis=basis):
v, t = vt_tuple
return _fvib(np.dot(np.stack(basis(v, t), axis=-1), coef), t)
return fit_func
[docs]
def get_poly_fit_func(basis):
def fit_func(vt_tuple, *coef, basis=basis):
return np.dot(np.stack(basis(*vt_tuple), axis=-1), coef)
return fit_func
[docs]
def get_fel_fit_func(basis):
raise NotImplementedError
def fit_func(vt_tuple, *coef, basis=basis):
return None
return fit_func
# def _v_basis(volume):
# """This function defines basis for V only, when T is fixed.
# This is used to fit to the surface (fitted using basis(T, V))
# on a mesh of fixed T."""
# one = 1. + volume*0.
# return np.array([one, volume, volume**2, volume**3])
# if weighted_fitting:
# # the error is given as 95% conf interval, divide by 1.96 to get SE
# fit_V_T_result = scipy.optimize.curve_fit(fit_func_for_curve_fit, (Fsurf[:, 0], Fsurf[:, 1]), Fsurf[:,2], p0=startcoef, sigma=Fsurf[:,3]/1.96, absolute_sigma=True)
#
#
# def fit_over_volumes_to_surface_fit(basis2, vminmax, tminmax, tstep, fitFunc, fitted_coef_V_T):
# volmin, volmax = vminmax
# tmin, tmax = tminmax
# volume_mesh = np.linspace(volmin, volmax, 101)
# i = 1
# deltaMax = 0
# # Least square fit for a linear function
# volume_basis_matrix = np.array([basis2(vol) for vol in volume_mesh])
# for T_current in range(tmin, tmax + tstep, tstep):
# values_on_mesh = np.array([fitFunc(vol, T_current, fitted_coef_V_T) for vol in volume_mesh])
# result = lstsq(volume_basis_matrix, values_on_mesh)
# fitted_coef_V = result[0]
# delta_max_diff = np.max(np.abs(np.matmul(volume_basis_matrix, fitted_coef_V) - values_on_mesh))
# if delta_max_diff > deltaMax:
# deltaMax = delta_max_diff
# diff_length = len(basis2(1)) + 1
# if diff_length > 0:
# for i_diff in range(diff_length):
# return deltaMax, volume_mesh