src/mdevaluate/extra/free_energy_landscape.py

@@ -47,21 +47,6 @@ 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,
@@ -128,14 +113,23 @@ def find_maxima(
     maxima_df = occupation_df.copy()
     maxima_df["maxima"] = None
     points = np.array(maxima_df[["x", "y", "z"]])
-    tree, points_pbc_index = _build_tree(points, box, radius, pore_geometry)
+    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]
     for i in range(len(maxima_df)):
         if maxima_df.loc[i, "maxima"] is not None:
             continue
-        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 = np.array(all_neighbors[i])
         neighbors = neighbors[neighbors != i]
         if len(neighbors) == 0:
             maxima_df.loc[i, "maxima"] = True
@@ -160,7 +154,16 @@ def _calc_energies(
     nodes: int = 8,
 ) -> NDArray:
     points = np.array(maxima_df[["x", "y", "z"]])
-    tree, points_pbc_index = _build_tree(points, box, bins[-1], pore_geometry)
+    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)
     maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]]
     maxima_occupations = np.array(maxima_df.loc[maxima_indices, "occupation"])
     num_of_neighbors = np.max(
@@ -184,7 +187,7 @@ def _calc_energies(
     all_occupied_bins_hist = []
     if distances.ndim == 1:
         current_distances = distances[1:][distances[1:] <= bins[-1]]
-        if points_pbc_index is None:
+        if np.all(np.diag(np.diag(box)) == box):
             current_indices = indices[1:][distances[1:] <= bins[-1]]
         else:
             current_indices = points_pbc_index[indices[1:][distances[1:] <= bins[-1]]]
@@ -198,7 +201,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 points_pbc_index is None:
+        if np.all(np.diag(np.diag(box)) == box):
             current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]]
         else:
             current_indices = points_pbc_index[