import logging
import warnings
from collections.abc import Callable
from dataclasses import dataclass
import numpy as np
import pandas as pd
import scipy.integrate
import scipy.optimize
import scipy.stats
from numpy.polynomial import Polynomial
from directupsampling.parallel import parlog
logger = logging.getLogger(__name__)
[docs]
@dataclass
class RawData:
"""Container for raw data used in thermodynamic integration fitting."""
lambda_: np.ndarray
energy: np.ndarray
error: np.ndarray
[docs]
class ThermoIntFit:
fit_types = ("linear", "cubic", "tan", "tan2", "glogit")
def __init__(
self,
lambdas: np.ndarray,
energies: np.ndarray,
errors=[],
) -> None:
self.raw_data = RawData(
np.array(lambdas, dtype=float),
np.array(energies, dtype=float),
np.array(errors, dtype=float),
)
self.fit_results = {}
self.fit()
@property
def best_fit(self):
df = pd.DataFrame.from_dict(self.fit_results, orient="index")
return df.loc[df["fit error"].idxmin()]
def __repr__(self) -> str:
return f"ThermoIntFit({self.__dict__!r})"
[docs]
def fit(self) -> None:
lambda_, energy, error = self.check_raw_data()
f_trapz = scipy.integrate.trapezoid(energy, x=lambda_)
f_simps = scipy.integrate.simpson(energy, x=lambda_)
for msg in [
f" (trapz: {f_trapz:.5f})",
f" (simpson: {f_simps:.5f})",
]:
parlog(msg, logger=logger, level="debug")
for fit_type in self.fit_types:
try:
fun = _get_fit(lambda_, energy, fit_type)
kwargs = {"limit": 1000, "epsabs": 1e-5}
res = scipy.integrate.quad(fun, 0, 1, **kwargs)
f_integrated = res[0]
except RuntimeError:
msg = f"`fit_type` '{fit_type}' could not be fitted"
parlog(msg, logger, level="warning")
continue
fit_err = _get_error_estimate(fun, lambda_, energy, NPARAMS_DICT[fit_type])
self.fit_results[fit_type] = {
"free energy": f_integrated,
"fit error": fit_err,
"fit": fun,
}
parlog(
f" {fit_type + ':':9s} {f_integrated:.5f} \u00b1 {fit_err:.5f}",
logger,
level="debug",
)
[docs]
def check_raw_data(self) -> tuple:
where_finite = np.isfinite(self.raw_data.energy)
if not np.all(where_finite):
msg = (
f"Raw data contains nonfinite values: "
f"{self.raw_data.energy[~where_finite]} for lambdas "
f"{self.raw_data.lambda_[~where_finite]}."
)
warnings.warn(msg)
energy = self.raw_data.energy[where_finite]
lambda_ = self.raw_data.lambda_[where_finite]
if len(self.raw_data.error) > 0:
error = self.raw_data.error[where_finite]
else:
error = []
return lambda_, energy, error
def _get_error_estimate(
fit_fun: Callable,
x: np.ndarray,
y: np.ndarray,
nparam: int,
error_method: str = "t_distribution",
) -> float:
"""Estimate error of fit_fun compared to data (x, y).
Parameters
----------
nparam : int
Number of parameters of the fit_fun.
Excludes non-finite values, consistent with the fitting.
Returns
-------
err : float
"""
where_finite = np.isfinite(y)
x = x[where_finite]
y = y[where_finite]
if len(x) < 6:
error_method = "mean_abs"
diff = fit_fun(x) - y
if error_method == "mean_abs":
err = np.average(np.abs(diff))
elif error_method == "t_distribution":
npts = len(x)
std_err = scipy.stats.sem(diff, ddof=nparam)
err = scipy.stats.t.interval(0.95, df=npts - nparam, loc=0, scale=std_err)[1]
return err
########## Fitting definitions ##########
def _get_fit(
lambda_: float,
energy: np.ndarray,
fit_type: str = "tan",
) -> Callable:
"""Get a function for fitting.
Parameters
----------
fit_type : str
"tan", "tan2", "logit", "glogit", "linear", "cubic" or an int determining
the order of a polynimial fit.
Returns
-------
Callable
Fitted function based on "fit_type".
Raises
------
RuntimeError
"""
if fit_type == "linear":
return Polynomial.fit(lambda_, energy, deg=1)
if fit_type == "cubic":
return Polynomial.fit(lambda_, energy, deg=3)
if isinstance(fit_type, int):
return Polynomial.fit(lambda_, energy, deg=fit_type)
if fit_type in {"tan2", "glogit"}:
prefit = {"tan2": "tan", "glogit": "logit"}
fun, param = _fit_custom_function(
CUSTOM_FUNCTION_DICT[prefit[fit_type]],
lambda_,
energy,
return_param=True,
)
fun_dict = CUSTOM_FUNCTION_DICT[fit_type]
fun_dict["p0"] = [*list(param), 0]
return _fit_custom_function(fun_dict, lambda_, energy)
if fit_type in CUSTOM_FUNCTION_DICT:
return _fit_custom_function(CUSTOM_FUNCTION_DICT[fit_type], lambda_, energy)
msg = f"fit_type {fit_type} not recognized"
raise RuntimeError(msg)
# # # # # # # # # # # # #
def _fit_custom_function(
function_dict: dict,
lambda_: float,
energy: np.ndarray,
return_param: bool = False,
) -> Callable:
"""Return function fitted to data.
Parameters
----------
function_dict : dict
"function", "bounds" and "p0" (initial guess of parameters).
The value for "function" should take the variables as
first argument and parameters as following arguments
(suited for scipy.optimize.curve_fit).
Returns
-------
Callable
"""
function, bounds, p0 = [function_dict[k] for k in ["function", "bounds", "p0"]]
p0[0] = np.mean(energy) # ini guess for y shift
param = scipy.optimize.curve_fit(function, lambda_, energy, p0=p0, bounds=bounds)[0]
def fitted_function(x):
return function(x, *param)
return (fitted_function, param) if return_param else fitted_function
# # # # # # # # # # # # #
NPARAMS_DICT = {
"linear": 2,
"cubic": 4,
"tan": 4,
"tan2": 5,
"logit": 5,
"glogit": 6,
}
[docs]
def tan(x, a0, a1, a2, a3):
"""Cotangens"""
# return a0 / np.tan(a1*np.pi*(x + a2)) + a3
return a0 + a1 / np.tan(np.pi * (a2 + x) * a3)
[docs]
def tan2(x, a0, a1, a2_over_a3_a4, a3, a4):
"""Cotangens(x**2 + ...)"""
# a2 = a2_over_a1_a4*(a1 + a4)
# return a0 / np.tan(np.pi*(a4*x**2 + a1*x + a2)) + a3
a2 = a2_over_a3_a4 * (a3 + a4)
return a0 + a1 / np.tan(np.pi * (a2 + a3 * x + a4 * x**2))
[docs]
def logit(x, a0, a1, a2, a3, a4):
"""Logit function
a0: y shift
a1: slope
a2: higher x shift (x of y = inf)
a3: lower x shift (-x of y = -inf)
a4: y span
"""
# return a0 * np.log(((a2 + a1)/(x + a1) - 1)/a4) + a3
return a0 + a1 * np.log(((a2 + a3) / (x + a3) - 1) / a4)
[docs]
def glogit(x, a0, a1, a2, a3, a4, a5):
"""Generalized logit function"""
# return a0 * np.log((((a2 + a1)/(x + a1))**(a5+1) - 1)/a4) + a3
return a0 + a1 * np.log((((a2 + a3) / (x + a3)) ** (a5 + 1) - 1) / a4)
CUSTOM_FUNCTION_DICT = {
# '''p0[0] should be y shift'''
"tan": {
"function": tan,
# 'bounds': ([0, 0, 0, -np.inf], [np.inf, 1, 1, np.inf]),
# 'p0': [200, 0.7, 0.1, -0.6]},
"bounds": ([-np.inf, 0, 0, 0], [np.inf, np.inf, 1, 1]),
"p0": [-0.6, 200, 0.1, 0.7],
},
"tan2": {
"function": tan2,
# 'bounds': ([0, 0, 0, -np.inf, -1], [np.inf, 1, 1, np.inf, 1]),
# 'p0': [200, 0.7, 0.1, -0.6, 0]},
"bounds": ([-np.inf, 0, 0, 0, -1], [np.inf, np.inf, 1, 1, 1]),
"p0": [-0.6, 200, 0.1, 0.7, 0],
},
"logit": {
"function": logit,
# 'bounds': ([0, 1e-5, 1+1e-5, -np.inf, 0],
"bounds": (
[-np.inf, 0, 1 + 1e-5, 1e-5, 0],
[np.inf, np.inf, np.inf, np.inf, np.inf],
),
# 'p0': [50, 1e-2, 1.2, -100, 1]},
"p0": [-100, 50, 1.2, 1e-2, 1],
},
"glogit": {
"function": glogit,
# 'bounds': ([0, 1e-5, 1+1e-5, -np.inf, 0, -1],
"bounds": (
[-np.inf, 0, 1 + 1e-5, 1e-5, 0, -1],
[np.inf, np.inf, np.inf, np.inf, np.inf, 100],
),
# 'p0': [50, 1e-2, 1.2, -100, 1, 0]},
"p0": [-100, 50, 1.2, 1e-2, 1, 0],
},
}
##################################################