src/mdevaluate/extra/free_energy_landscape.py

@@ -4,6 +4,7 @@ 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
@@ -48,7 +49,7 @@ def _pbc_points_reduced(
 
 def _build_tree(points, box, r_max, pore_geometry):
     if np.all(np.diag(np.diag(box)) == box):
-        tree = KDTree(points % box, boxsize=box)
+        tree = KDTree(points, boxsize=box)
         points_pbc_index = None
     else:
         points_pbc, points_pbc_index = _pbc_points_reduced(
@@ -78,7 +79,8 @@ 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
-    indices = np.array_split(frame_indices, nodes)
+    size = math.ceil(len(frame_indices) / nodes)
+    indices = [frame_indices[i : i + size] for i in range(0, len(frame_indices), size)]
     pool = mp.Pool(nodes)
     results = pool.map(
         partial(_calc_histogram, trajectory=trajectory, bins=bins), indices
@@ -272,11 +274,7 @@ def distance_resolved_energies(
 
 
 def find_energy_maxima(
-    energy_df: pd.DataFrame,
-    r_min: float,
-    r_max: float,
-    r_eval: float = None,
-    degree: int = 2,
+    energy_df: pd.DataFrame, r_min: float, r_max: float
 ) -> pd.DataFrame:
     distances = []
     energies = []
@@ -285,9 +283,6 @@ 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=degree)
-        if r_eval is None:
-            energies.append(np.max(p3(np.linspace(r_min, r_max, 1000))))
-        else:
-            energies.append(p3(r_eval))
+        p3 = Poly.fit(x[mask], y[mask], deg=2)
+        energies.append(np.max(p3(np.linspace(r_min, r_max, 1000))))
     return pd.DataFrame({"d": distances, "energy": energies})