From 0968ff214c2854c5c9d6013c107f31401cd1a26c Mon Sep 17 00:00:00 2001 From: Sebastian Kloth Date: Fri, 5 Jan 2024 16:04:14 +0100 Subject: [PATCH] Completed new functions for free energy landscape calculations --- src/mdevaluate/extra/free_energy_landscape.py | 86 +++++++++++++++---- 1 file changed, 70 insertions(+), 16 deletions(-) diff --git a/src/mdevaluate/extra/free_energy_landscape.py b/src/mdevaluate/extra/free_energy_landscape.py index 673b917..76d4568 100644 --- a/src/mdevaluate/extra/free_energy_landscape.py +++ b/src/mdevaluate/extra/free_energy_landscape.py @@ -2,9 +2,10 @@ from functools import partial import os.path import numpy as np +from numpy.typing import ArrayLike, NDArray +from numpy.polynomial.polynomial import Polynomial as Poly import math import scipy -from numpy.typing import ArrayLike, NDArray from scipy.spatial import KDTree import cmath import pandas as pd @@ -103,7 +104,13 @@ def find_maxima( return maxima_df -def calc_energies(maxima_indices, maxima_df, bins, box): +def _calc_energies( + maxima_indices: ArrayLike, + maxima_df: pd.DataFrame, + bins: ArrayLike, + box: NDArray, + T: float, +) -> NDArray: points = np.array(maxima_df[["x", "y", "z"]]) tree = KDTree(points, boxsize=box) maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]] @@ -119,8 +126,9 @@ def calc_energies(maxima_indices, maxima_df, bins, box): if distances.ndim == 1: current_distances = distances[1:][distances[1:] <= bins[-1]] current_indices = indices[1:][distances[1:] <= bins[-1]] - energy = -np.log( - maxima_df.loc[current_indices, "occupation"] / maxima_occupations + energy = ( + -np.log(maxima_df.loc[current_indices, "occupation"] / maxima_occupations) + * T ) energy_hist = np.histogram(current_distances, bins=bins, weights=energy)[0] occupied_bins_hist = np.histogram(current_distances, bins=bins)[0] @@ -130,8 +138,9 @@ def calc_energies(maxima_indices, maxima_df, bins, box): current_distances = distances[i, 1:][distances[i, 1:] <= bins[-1]] current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]] - energy = -np.log( - maxima_df.loc[current_indices, "occupation"] / maxima_occupation + energy = ( + -np.log(maxima_df.loc[current_indices, "occupation"] / maxima_occupation) + * T ) energy_hist = np.histogram(current_distances, bins=bins, weights=energy)[0] occupied_bins_hist = np.histogram(current_distances, bins=bins)[0] @@ -141,20 +150,65 @@ def calc_energies(maxima_indices, maxima_df, bins, box): return result -def add_distances(occupation_df, pore_geometry, origin): - new_df = occupation_df.copy() - if pore_geometry is "cylindrical": - new_df["distance"] = ( - (new_df["x"] - origin[0]) ** 2 + (new_df["y"] - origin[1]) ** 2 +def add_distances( + occupation_df: pd.DataFrame, pore_geometry: str, origin: ArrayLike +) -> pd.DataFrame: + distance_df = occupation_df.copy() + if pore_geometry == "cylindrical": + distance_df["distance"] = ( + (distance_df["x"] - origin[0]) ** 2 + (distance_df["y"] - origin[1]) ** 2 ) ** (1 / 2) - elif pore_geometry is "slit": - new_df["distance"] = np.abs(new_df["z"] - origin[2]) + elif pore_geometry == "slit": + distance_df["distance"] = np.abs(distance_df["z"] - origin[2]) else: raise ValueError( f"pore_geometry is {pore_geometry}, should either be " f"'cylindrical' or 'slit'" ) - return new_df + return distance_df + + +def distance_resolved_energies( + maxima_df: pd.DataFrame, + distance_bins: ArrayLike, + r_bins: ArrayLike, + box: NDArray, + T: float, +) -> pd.DataFrame: + results = [] + for i in range(len(distance_bins) - 1): + maxima_indices = np.array( + maxima_df.index[ + (maxima_df["distance"] >= distance_bins[i]) + * (maxima_df["distance"] < distance_bins[i + 1]) + * (maxima_df["maxima"] == True) + ] + ) + results.append(_calc_energies(maxima_indices, maxima_df, r_bins, box, T)) + + distances = (distance_bins[:-1] + distance_bins[1:]) / 2 + radii = (r_bins[:-1] + r_bins[1:]) / 2 + d = np.array([d for d in distances for r in radii]) + r = np.array([r for d in distances for r in radii]) + result = np.array(results).flatten() + return pd.DataFrame({"d": d, "r": r, "energy": result}) + + +def find_energy_maxima( + energy_df: pd.DataFrame, r_min: float, r_max: float +) -> pd.DataFrame: + distances = [] + energies = [] + for d, data_d in energy_df.groupby("d"): + distances.append(d) + x = np.array(data_d["r"]) + y = np.array(data_d["energy"]) + mask = (x >= r_min) * (x <= r_max) + print(x[mask], y[mask]) + p3 = Poly.fit(x[mask], y[mask], deg=2) + print(p3) + energies.append(np.max(p3(np.linspace(r_min, r_max, 1000)))) + return pd.DataFrame({"d": distances, "energy": energies}) def get_fel( @@ -214,8 +268,8 @@ def get_fel( np.save(f"{path}/radiiBins", bins) energy_differences = np.array(calculate_energy_difference(data, bins, temperature)) - r = bins[1:] - (bins[1] - bins[0]) / 2 - return r, energy_differences + d = np.linspace(radiusmin, radiusmax, len(energy_differences)) + return d, energy_differences def fill_bins(traj, path, edge=0.05):