Added Robins adjustments for FEL

src/mdevaluate/extra/free_energy_landscape.py

@@ -4,7 +4,6 @@ from typing import Optional
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 from numpy.polynomial.polynomial import Polynomial as Poly
-import math
 from scipy.spatial import KDTree
 import pandas as pd
 import multiprocessing as mp
@@ -47,6 +46,21 @@ def _pbc_points_reduced(
     return coordinates_pbc, indices
 
 
+def _build_tree(points, box, r_max, pore_geometry):
+    if np.all(np.diag(np.diag(box)) == box):
+        tree = KDTree(points % box, boxsize=box)
+        points_pbc_index = None
+    else:
+        points_pbc, points_pbc_index = _pbc_points_reduced(
+            points,
+            pore_geometry,
+            box,
+            thickness=r_max + 0.01,
+        )
+        tree = KDTree(points_pbc)
+    return tree, points_pbc_index
+
+
 def occupation_matrix(
     trajectory: Coordinates,
     edge_length: float = 0.05,
@@ -64,11 +78,7 @@ def occupation_matrix(
     z_bins = np.arange(0, box[2][2] + edge_length, edge_length)
     bins = [x_bins, y_bins, z_bins]
     # Trajectory is split for parallel computing
-    size = math.ceil(len(frame_indices) / nodes)
-    indices = [
-        np.arange(len(frame_indices))[i : i + size]
-        for i in range(0, len(frame_indices), size)
-    ]
+    indices = np.array_split(frame_indices, nodes)
     pool = mp.Pool(nodes)
     results = pool.map(
         partial(_calc_histogram, trajectory=trajectory, bins=bins), indices
@@ -113,23 +123,14 @@ def find_maxima(
     maxima_df = occupation_df.copy()
     maxima_df["maxima"] = None
     points = np.array(maxima_df[["x", "y", "z"]])
-    if np.all(np.diag(np.diag(box)) == box):
-        tree = KDTree(points, boxsize=box)
-        all_neighbors = tree.query_ball_point(points, radius)
-    else:
-        points_pbc, points_pbc_index = _pbc_points_reduced(
-            points,
-            pore_geometry,
-            box,
-            thickness=radius + 0.01,
-        )
-        tree = KDTree(points_pbc)
-        all_neighbors = tree.query_ball_point(points, radius)
-        all_neighbors = points_pbc_index[all_neighbors]
+    tree, points_pbc_index = _build_tree(points, box, radius, pore_geometry)
     for i in range(len(maxima_df)):
         if maxima_df.loc[i, "maxima"] is not None:
             continue
-        neighbors = np.array(all_neighbors[i])
+        maxima_pos = maxima_df.loc[i, ["x", "y", "z"]]
+        neighbors = np.array(tree.query_ball_point(maxima_pos, radius))
+        if points_pbc_index is not None:
+            neighbors = points_pbc_index[neighbors]
         neighbors = neighbors[neighbors != i]
         if len(neighbors) == 0:
             maxima_df.loc[i, "maxima"] = True
@@ -154,16 +155,7 @@ def _calc_energies(
     nodes: int = 8,
 ) -> NDArray:
     points = np.array(maxima_df[["x", "y", "z"]])
-    if np.all(np.diag(np.diag(box)) == box):
-        tree = KDTree(points, boxsize=box)
-    else:
-        points_pbc, points_pbc_index = _pbc_points_reduced(
-            points,
-            pore_geometry,
-            box,
-            thickness=bins[-1] + 0.01,
-        )
-        tree = KDTree(points_pbc)
+    tree, points_pbc_index = _build_tree(points, box, bins[-1], pore_geometry)
     maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]]
     maxima_occupations = np.array(maxima_df.loc[maxima_indices, "occupation"])
     num_of_neighbors = np.max(
@@ -187,7 +179,7 @@ def _calc_energies(
     all_occupied_bins_hist = []
     if distances.ndim == 1:
         current_distances = distances[1:][distances[1:] <= bins[-1]]
-        if np.all(np.diag(np.diag(box)) == box):
+        if points_pbc_index is None:
             current_indices = indices[1:][distances[1:] <= bins[-1]]
         else:
             current_indices = points_pbc_index[indices[1:][distances[1:] <= bins[-1]]]
@@ -201,7 +193,7 @@ def _calc_energies(
         return result
     for i, maxima_occupation in enumerate(maxima_occupations):
         current_distances = distances[i, 1:][distances[i, 1:] <= bins[-1]]
-        if np.all(np.diag(np.diag(box)) == box):
+        if points_pbc_index is None:
             current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]]
         else:
             current_indices = points_pbc_index[
@@ -280,7 +272,11 @@ def distance_resolved_energies(
 
 
 def find_energy_maxima(
-    energy_df: pd.DataFrame, r_min: float, r_max: float
+    energy_df: pd.DataFrame,
+    r_min: float,
+    r_max: float,
+    r_eval: float = None,
+    degree: int = 2,
 ) -> pd.DataFrame:
     distances = []
     energies = []
@@ -289,6 +285,9 @@ def find_energy_maxima(
         x = np.array(data_d["r"])
         y = np.array(data_d["energy"])
         mask = (x >= r_min) * (x <= r_max)
-        p3 = Poly.fit(x[mask], y[mask], deg=2)
-        energies.append(np.max(p3(np.linspace(r_min, r_max, 1000))))
+        p3 = Poly.fit(x[mask], y[mask], deg=degree)
+        if r_eval is None:
+            energies.append(np.max(p3(np.linspace(r_min, r_max, 1000))))
+        else:
+            energies.append(p3(r_eval))
     return pd.DataFrame({"d": distances, "energy": energies})

test/test_free_energy_landscape.py

@@ -13,19 +13,7 @@ def trajectory(request):
 
 
 def test_get_fel(trajectory):
-    test_array = np.array(
-        [
-            184.4909136,
-            233.79320471,
-            223.12003988,
-            228.49746397,
-            200.5626769,
-            212.82484221,
-            165.10818396,
-            170.74123681,
-            175.86672931,
-        ]
-    )
+    test_array = np.array([210., 214., 209., 192., 200., 193., 230., 218., 266.])
 
     OW = trajectory.subset(atom_name="OW")
     box = trajectory[0].box