Completed new functions for free energy landscape calculations

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):