Added Robins adjustments for FEL

# Conflicts:
#	src/mdevaluate/coordinates.py

src/mdevaluate/coordinates.py

@@ -218,7 +218,7 @@ class Coordinates:
             self.get_frame.clear_cache()
 
     def __iter__(self):
-        for i in range(len(self))[self._slice]:
+        for i in range(len(self.frames))[self._slice]:
             yield self[i]
 
     @singledispatchmethod
@@ -232,7 +232,7 @@ class Coordinates:
         return sliced
 
     def __len__(self):
-        return len(self.frames)
+        return len(self.frames[self._slice])
 
     def __checksum__(self):
         return checksum(self.frames, self.atom_filter, self._slice, self.mode)

src/mdevaluate/extra/chill.py

@@ -11,7 +11,7 @@ from mdevaluate.coordinates import CoordinateFrame, Coordinates
 from mdevaluate.pbc import pbc_points
 
 
-def a_ij(atoms: ArrayLike, N: int = 4, l: int = 3) -> tuple[NDArray, NDArray]:
+def calc_aij(atoms: ArrayLike, N: int = 4, l: int = 3) -> tuple[NDArray, NDArray]:
     tree = KDTree(atoms)
 
     dist, indices = tree.query(atoms, N + 1)
@@ -84,18 +84,18 @@ def count_ice_types(iceTypes: NDArray) -> NDArray:
 
 
 def selector_ice(
-    start_frame: CoordinateFrame,
-    traj: Coordinates,
+    oxygen_atoms_water: CoordinateFrame,
     chosen_ice_types: ArrayLike,
     combined: bool = True,
+    next_neighbor_distance: float = 0.35,
 ) -> NDArray:
-    oxygen = traj.subset(atom_name="OW")
-    atoms = oxygen[start_frame.step]
-    atoms = atoms % np.diag(atoms.box)
-    atoms_PBC = pbc_points(atoms, thickness=1)
-    aij, indices = a_ij(atoms_PBC)
+    atoms = oxygen_atoms_water
+    atoms_PBC = pbc_points(atoms, thickness=next_neighbor_distance * 2.2)
+    aij, indices = calc_aij(atoms_PBC)
     tree = KDTree(atoms_PBC)
-    neighbors = tree.query_ball_point(atoms_PBC, 0.35, return_length=True)
+    neighbors = tree.query_ball_point(
+        atoms_PBC, next_neighbor_distance, return_length=True
+    ) - 1
     index_SOL = atoms_PBC.tolist().index(atoms[0].tolist())
     index_SOL = np.arange(index_SOL, index_SOL + len(atoms))
     ice_Types = classify_ice(aij, indices, neighbors, index_SOL)
@@ -117,9 +117,9 @@ def selector_ice(
 def ice_types(trajectory: Coordinates, segments: int = 10000) -> pd.DataFrame:
     def ice_types_distribution(frame: CoordinateFrame, selector: Callable) -> NDArray:
         atoms_PBC = pbc_points(frame, thickness=1)
-        aij, indices = a_ij(atoms_PBC)
+        aij, indices = calc_aij(atoms_PBC)
         tree = KDTree(atoms_PBC)
-        neighbors = tree.query_ball_point(atoms_PBC, 0.35, return_length=True)
+        neighbors = tree.query_ball_point(atoms_PBC, 0.35, return_length=True) - 1
         index = selector(frame, atoms_PBC)
         ice_types_data = classify_ice(aij, indices, neighbors, index)
         ice_parts_data = count_ice_types(ice_types_data)
@@ -161,8 +161,8 @@ def ice_types(trajectory: Coordinates, segments: int = 10000) -> pd.DataFrame:
 def ice_clusters_traj(
     traj: Coordinates, segments: int = 10000, skip: float = 0.1
 ) -> list:
-    def ice_clusters(frame: CoordinateFrame, traj: Coordinates) -> Tuple[float, list]:
-        selection = selector_ice(frame, traj, [0, 1, 2])
+    def ice_clusters(frame: CoordinateFrame) -> Tuple[float, list]:
+        selection = selector_ice(frame, [0, 1, 2])
         if len(selection) == 0:
             return frame.time, []
         else:
@@ -194,6 +194,6 @@ def ice_clusters_traj(
         np.int_(np.linspace(len(traj) * skip, len(traj) - 1, num=segments))
     )
     all_clusters = [
-        ice_clusters(traj[frame_index], traj) for frame_index in frame_indices
+        ice_clusters(traj[frame_index]) for frame_index in frame_indices
     ]
     return all_clusters

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,48 +13,24 @@ def trajectory(request):
 
 
 def test_get_fel(trajectory):
-    test_array = np.array(
-        [
-            174.46253634,
-            174.60905476,
-            178.57658092,
-            182.43001192,
-            180.57916378,
-            176.49886217,
-            178.96018547,
-            181.13561782,
-            178.31026314,
-            176.08903996,
-            180.71215345,
-            181.59703135,
-            180.34329368,
-            187.02474488,
-            197.99167477,
-            214.05788031,
-            245.58571282,
-            287.52457507,
-            331.53492965,
-        ]
-    )
+    test_array = np.array([210., 214., 209., 192., 200., 193., 230., 218., 266.])
 
     OW = trajectory.subset(atom_name="OW")
-
     box = trajectory[0].box
     box_voxels = (np.diag(box) // [0.05, 0.05, 0.05] + [1, 1, 1]) * [0.05, 0.05, 0.05]
-    occupation_matrix = fel.occupation_matrix(OW, skip=0, segments=1000)
+    occupation_matrix = fel.occupation_matrix(OW, skip=0, segments=10)
     radius_maxima = 0.05 * 3 ** (1 / 2) + 0.05 / 100
     maxima_matrix = fel.find_maxima(
         occupation_matrix,
         box=box_voxels,
         radius=radius_maxima,
-        pore_geometry="cylindrical"
+        pore_geometry="cylindrical",
     )
     maxima_matrix = fel.add_distances(maxima_matrix, "cylindrical", np.diag(box) / 2)
-    r_bins = np.arange(0, 1, 0.02)
-    distance_bins = np.arange(0.05, 2.05, 0.1)
+    r_bins = np.arange(0, 0.5, 0.02)
+    distance_bins = np.arange(1.8, 1.9, 0.01)
     energy_df = fel.distance_resolved_energies(
         maxima_matrix, distance_bins, r_bins, box, "cylindrical", 225
     )
     result = fel.find_energy_maxima(energy_df, r_min=0.05, r_max=0.15)
-
     assert (np.round(np.array(result["energy"])) == np.round(test_array)).all()