From c89cead81c67c8c24856f8fc1ba71178d513afff Mon Sep 17 00:00:00 2001
From: Sebastian Kloth <sebastian.kloth@pkm.tu-darmstadt.de>
Date: Tue, 5 Mar 2024 13:57:55 +0100
Subject: [PATCH] Moved KDTree build up in its own function and made
 find_maxima query for each point separately.

---
 src/mdevaluate/extra/free_energy_landscape.py | 49 +++++++++----------
 1 file changed, 23 insertions(+), 26 deletions(-)

diff --git a/src/mdevaluate/extra/free_energy_landscape.py b/src/mdevaluate/extra/free_energy_landscape.py
index 5190d5c..85a8c25 100644
--- a/src/mdevaluate/extra/free_energy_landscape.py
+++ b/src/mdevaluate/extra/free_energy_landscape.py
@@ -47,6 +47,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, 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,
@@ -113,23 +128,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 +160,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 +184,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 +198,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[