Added new function for energy calculations

src/mdevaluate/extra/free_energy_landscape.py

@@ -4,15 +4,24 @@ import os.path
 import numpy as np
 import math
 import scipy
+from numpy.typing import ArrayLike, NDArray
 from scipy.spatial import KDTree
 import cmath
 import pandas as pd
 import multiprocessing as mp
 
+from ..coordinates import Coordinates
+
 VALID_GEOMETRY = {"cylindrical", "slab"}
 
 
-def occupation_matrix(trajectory, edge_length=0.05, segments=1000, skip=0.1, nodes=8):
+def occupation_matrix(
+    trajectory: Coordinates,
+    edge_length: float = 0.05,
+    segments: int = 1000,
+    skip: float = 0.1,
+    nodes: int = 8,
+) -> pd.DataFrame:
     frame_indices = np.unique(
         np.int_(np.linspace(len(trajectory) * skip, len(trajectory) - 1, num=segments))
     )
@@ -48,7 +57,9 @@ def occupation_matrix(trajectory, edge_length=0.05, segments=1000, skip=0.1, nod
     return occupation_df
 
 
-def _calc_histogram(numberlist, trajectory, bins):
+def _calc_histogram(
+    numberlist: ArrayLike, trajectory: Coordinates, bins: ArrayLike
+) -> NDArray:
     matbin = None
     for index in range(0, len(numberlist), 1000):
         try:
@@ -64,7 +75,9 @@ def _calc_histogram(numberlist, trajectory, bins):
     return matbin
 
 
-def find_maxima(occupation_df, box, edge_length=0.05):
+def find_maxima(
+    occupation_df: pd.DataFrame, box: ArrayLike, edge_length: float = 0.05
+) -> pd.DataFrame:
     maxima_df = occupation_df.copy()
     maxima_df["maxima"] = None
     points = np.array(maxima_df[["x", "y", "z"]])
@@ -80,8 +93,8 @@ def find_maxima(occupation_df, box, edge_length=0.05):
         if len(neighbors) == 0:
             maxima_df.loc[i, "maxima"] = True
         elif (
-                maxima_df.loc[neighbors, "occupation"].max()
+            maxima_df.loc[neighbors, "occupation"].max()
-                < maxima_df.loc[i, "occupation"]
+            < maxima_df.loc[i, "occupation"]
         ):
             maxima_df.loc[neighbors, "maxima"] = False
             maxima_df.loc[i, "maxima"] = True
@@ -90,6 +103,44 @@ def find_maxima(occupation_df, box, edge_length=0.05):
     return maxima_df
 
 
+def calc_energies(maxima_indices, maxima_df, bins):
+    points = np.array(maxima_df[["x", "y", "z"]])
+    tree = KDTree(points, boxsize=box)
+    maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]]
+    maxima_occupations = np.array(maxima_df.loc[maxima_indices, "occupation"])
+    num_of_neighbors = np.max(
+        tree.query_ball_point(maxima, bins[-1], return_length=True)
+    )
+    distances, indices = tree.query(
+        maxima, k=num_of_neighbors, distance_upper_bound=bins[-1]
+    )
+    all_energy_hist = []
+    all_occupied_bins_hist = []
+    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_hist = np.histogram(current_distances, bins=bins, weights=energy)[0]
+        occupied_bins_hist = np.histogram(current_distances, bins=bins)[0]
+        result = energy_hist / occupied_bins_hist
+        return r, result
+    for i, maxima_occupation in enumerate(maxima_occupations):
+        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_hist = np.histogram(current_distances, bins=bins, weights=energy)[0]
+        occupied_bins_hist = np.histogram(current_distances, bins=bins)[0]
+        all_energy_hist.append(energy_hist)
+        all_occupied_bins_hist.append(occupied_bins_hist)
+    result = np.sum(all_energy_hist, axis=0) / np.sum(all_occupied_bins_hist, axis=0)
+    return result
+
+
 def get_fel(
     traj,
     path,
@@ -309,7 +360,7 @@ def sphere_quotient(
 
     # Distances between maxima and other cubes in the system
     coordlist = coordlist.reshape(-1, 3)
-    numOfNeigbour = tree.query_ball_point(maxima[0], unitdist, return_length=True)
+    numOfNeigbour = np.max(tree.query_ball_point(maxima, unitdist, return_length=True))
     d, neighbourlist = tree.query(maxima_masked, k=numOfNeigbour, workers=-1)
 
     i = 0
@@ -365,7 +416,7 @@ def sphere_quotient_slab(
     maxima_masked = np.array(maxima)[mask]
 
     coordlist = coordlist.reshape(-1, 3)
-    numOfNeigbour = tree.query_ball_point(maxima[0], unitdist, return_length=True)
+    numOfNeigbour = np.max(tree.query_ball_point(maxima, unitdist, return_length=True))
     d, neighbourlist = tree.query(maxima_masked, k=numOfNeigbour, workers=-1)
 
     i = 0