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

@@ -1,9 +1,9 @@
 from functools import partial
+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
@@ -11,6 +11,56 @@ import multiprocessing as mp
 from ..coordinates import Coordinates
 
 
+def _pbc_points_reduced(
+    coordinates: ArrayLike,
+    pore_geometry: str,
+    box: Optional[NDArray] = None,
+    thickness: Optional[float] = None,
+) -> tuple[NDArray, NDArray]:
+    if box is None:
+        box = coordinates.box
+    if pore_geometry == "cylindrical":
+        grid = np.array([[i, j, k] for k in [-1, 0, 1] for j in [0] for i in [0]])
+        indices = np.tile(np.arange(len(coordinates)), 3)
+    elif pore_geometry == "slit":
+        grid = np.array(
+            [[i, j, k] for k in [0] for j in [1, 0, -1] for i in [-1, 0, 1]]
+        )
+        indices = np.tile(np.arange(len(coordinates)), 9)
+    else:
+        raise ValueError(
+            f"pore_geometry is {pore_geometry}, should either be "
+            f"'cylindrical' or 'slit'"
+        )
+    coordinates_pbc = np.concatenate([coordinates + v @ box for v in grid], axis=0)
+    size = np.diag(box)
+
+    if thickness is not None:
+        mask = np.all(coordinates_pbc > -thickness, axis=1)
+        coordinates_pbc = coordinates_pbc[mask]
+        indices = indices[mask]
+        mask = np.all(coordinates_pbc < size + thickness, axis=1)
+        coordinates_pbc = coordinates_pbc[mask]
+        indices = indices[mask]
+
+    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,
@@ -28,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
@@ -72,19 +118,19 @@ def _calc_histogram(
 
 
 def find_maxima(
-    occupation_df: pd.DataFrame, box: ArrayLike, edge_length: float = 0.05
+    occupation_df: pd.DataFrame, box: ArrayLike, radius: float, pore_geometry: str
 ) -> pd.DataFrame:
     maxima_df = occupation_df.copy()
     maxima_df["maxima"] = None
     points = np.array(maxima_df[["x", "y", "z"]])
-    tree = KDTree(points, boxsize=box)
-    all_neighbors = tree.query_ball_point(
-        points, edge_length * 3 ** (1 / 2) + edge_length / 100
-    )
+    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
@@ -104,23 +150,39 @@ def _calc_energies(
     maxima_df: pd.DataFrame,
     bins: ArrayLike,
     box: NDArray,
+    pore_geometry: str,
     T: float,
+    nodes: int = 8,
 ) -> NDArray:
     points = np.array(maxima_df[["x", "y", "z"]])
-    tree = KDTree(points, boxsize=box)
+    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(
         tree.query_ball_point(maxima, bins[-1], return_length=True)
     )
-    distances, indices = tree.query(
-        maxima, k=num_of_neighbors, distance_upper_bound=bins[-1]
-    )
+    split_maxima = []
+    for i in range(0, len(maxima), 1000):
+        split_maxima.append(maxima[i : i + 1000])
+
+    distances = []
+    indices = []
+    for maxima in split_maxima:
+        distances_step, indices_step = tree.query(
+            maxima, k=num_of_neighbors, distance_upper_bound=bins[-1], workers=nodes
+        )
+        distances.append(distances_step)
+        indices.append(indices_step)
+    distances = np.concatenate(distances)
+    indices = np.concatenate(indices)
     all_energy_hist = []
     all_occupied_bins_hist = []
     if distances.ndim == 1:
         current_distances = distances[1:][distances[1:] <= bins[-1]]
-        current_indices = indices[1:][distances[1:] <= bins[-1]]
+        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]]]
         energy = (
             -np.log(maxima_df.loc[current_indices, "occupation"] / maxima_occupations)
             * T
@@ -131,8 +193,12 @@ def _calc_energies(
         return result
     for i, maxima_occupation in enumerate(maxima_occupations):
         current_distances = distances[i, 1:][distances[i, 1:] <= bins[-1]]
-        current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]]
-
+        if points_pbc_index is None:
+            current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]]
+        else:
+            current_indices = points_pbc_index[
+                indices[i, 1:][distances[i, 1:] <= bins[-1]]
+            ]
         energy = (
             -np.log(maxima_df.loc[current_indices, "occupation"] / maxima_occupation)
             * T
@@ -168,9 +234,12 @@ def distance_resolved_energies(
     distance_bins: ArrayLike,
     r_bins: ArrayLike,
     box: NDArray,
+    pore_geometry: str,
     temperature: float,
+    nodes: int = 8,
 ) -> pd.DataFrame:
     results = []
+    distances = []
     for i in range(len(distance_bins) - 1):
         maxima_indices = np.array(
             maxima_df.index[
@@ -179,11 +248,22 @@ def distance_resolved_energies(
                 * (maxima_df["maxima"] == True)
             ]
         )
-        results.append(
-            _calc_energies(maxima_indices, maxima_df, r_bins, box, temperature)
-        )
+        try:
+            results.append(
+                _calc_energies(
+                    maxima_indices,
+                    maxima_df,
+                    r_bins,
+                    box,
+                    pore_geometry,
+                    temperature,
+                    nodes,
+                )
+            )
+            distances.append((distance_bins[i] + distance_bins[i + 1]) / 2)
+        except ValueError:
+            pass
 
-    distances = (distance_bins[:-1] + distance_bins[1:]) / 2
     radii = (r_bins[:-1] + r_bins[1:]) / 2
     d = np.array([d for d in distances for r in radii])
     r = np.array([r for d in distances for r in radii])
@@ -192,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 = []
@@ -201,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})

src/mdevaluate/extra/water.py

@@ -265,7 +265,7 @@ def LSI(
     )
     distributions = np.array(
         [
-            LSI_distribution(trajectory[frame_index], trajectory, bins, selector=None)
+            LSI_distribution(trajectory[frame_index], bins, selector=None)
             for frame_index in frame_indices
         ]
     )

test/test_free_energy_landscape.py

@@ -13,42 +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 = np.diag(trajectory[0].box)
-    box_voxels = (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)
-    maxima_matrix = fel.find_maxima(occupation_matrix, box=box_voxels, edge_length=0.05)
-    maxima_matrix = fel.add_distances(maxima_matrix, "cylindrical", box / 2)
-    r_bins = np.arange(0, 2, 0.02)
-    distance_bins = np.arange(0.05, 2.05, 0.1)
+    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=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",
+    )
+    maxima_matrix = fel.add_distances(maxima_matrix, "cylindrical", np.diag(box) / 2)
+    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, 225
+        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()